Merge pull request #5 from sjtubblythe/dev

fix issue when run in GPU environment
CSAILVision · May 6, 2019 · 2163454 · 2163454
2 parents bb2d35f + fd5a997
commit 2163454
Showing 1 changed file with 1 addition and 2 deletions.
diff --git 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
-