diff --git a/deepfinder/inference.py b/deepfinder/inference.py
index 535d2b558d29c904337e7b925ce612b957dffdec..d8624bdcfdebccb01ed3c597c125135ca88bba62 100644
--- a/deepfinder/inference.py
+++ b/deepfinder/inference.py
@@ -130,7 +130,7 @@ class Segment(core.DeepFinder):
         dataArray = (dataArray[:] - np.mean(dataArray[:])) / np.std(dataArray[:])  # normalize
         dataArray = np.pad(dataArray, self.pcrop, mode='constant')  # zeropad
         dim = dataArray.shape
-        #dataArray = np.reshape(dataArray, (1, dim[0], dim[1], dim[2], 1))  # reshape for keras [batch,x,y,z,channel]
+        dataArray = np.reshape(dataArray, (1, dim[0], dim[1], dim[2], 1))  # reshape for keras [batch,x,y,z,channel]
 
         pred = self.net.predict(dataArray, batch_size=1)
         predArray = pred[0, :, :, :, :]