diff --git a/src/msckf_vio.cpp b/src/msckf_vio.cpp
index 02b2879..330bd54 100644
--- a/src/msckf_vio.cpp
+++ b/src/msckf_vio.cpp
@@ -451,9 +451,13 @@ void MsckfVio::imageCallback(
   //cout << "1" << endl;
   // Augment the state vector.
   start_time = ros::Time::now();
-  //truePhotometricStateAugmentation(feature_msg->header.stamp.toSec());
-  PhotometricStateAugmentation(feature_msg->header.stamp.toSec());
-
+  if(PHOTOMETRIC)
+  {
+    truePhotometricStateAugmentation(feature_msg->header.stamp.toSec());
+    PhotometricStateAugmentation(feature_msg->header.stamp.toSec());
+  }
+  else
+    stateAugmentation(feature_msg->header.stamp.toSec());
   double state_augmentation_time = (
       ros::Time::now()-start_time).toSec();
 
@@ -1702,7 +1706,9 @@ void MsckfVio::measurementUpdate(const MatrixXd& H, const VectorXd& r) {
   // complexity as in Equation (28), (29).
   MatrixXd H_thin;
   VectorXd r_thin;
-  int augmentationSize = 7;
+  int augmentationSize = 6;
+  if(PHOTOMETRIC)
+    augmentationSize = 7;
 
   // QR decomposition to make stuff faster
   if (H.rows() > H.cols()) {
@@ -1823,7 +1829,9 @@ void MsckfVio::removeLostFeatures() {
   // Remove the features that lost track.
   // BTW, find the size the final Jacobian matrix and residual vector.
   int jacobian_row_size = 0;
-  int augmentationSize = 7;
+  int augmentationSize = 6;
+  if(PHOTOMETRIC)
+    augmentationSize = 7;
 
 
   vector<FeatureIDType> invalid_feature_ids(0);
@@ -2127,7 +2135,9 @@ void MsckfVio::pruneCamStateBuffer() {
 
   // Find the size of the Jacobian matrix.
   int jacobian_row_size = 0;
-    int augmentationSize = 7;
+    int augmentationSize = 6;
+    if(PHOTOMETRIC)
+      augmentationSize = 7;
 
   for (auto& item : map_server) {
     auto& feature = item.second;