diff --git a/feature_operation.py b/feature_operation.py
index 5890956..af6613b 100644
--- a/feature_operation.py
+++ b/feature_operation.py
@@ -65,7 +65,7 @@ def feature_extraction(self, model=None, memmap=True):
                 input = input.cuda()
             input_var = V(input,volatile=True)
             logit = model.forward(input_var)
-            while np.isnan(logit.data.max()):
+            while np.isnan(logit.data.cpu().max()):
                 print("nan") #which I have no idea why it will happen
                 del features_blobs[:]
                 logit = model.forward(input_var)
@@ -263,4 +263,3 @@ def tally(self, features, threshold, savepath=''):
                 for i in range(len(ordering)):
                     writer.writerow(rets[i])
         return rets
-