From fd5a9972a323b9dd485aa1ec42b86c39eaa70dfc Mon Sep 17 00:00:00 2001
From: Gaoyuan-Zhang <gaoyuan.zhang@ibm.com>
Date: Fri, 11 Jan 2019 17:25:18 -0500
Subject: [PATCH] fix issue when run in GPU environment

---
 feature_operation.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

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
-