testing if cuda works

bayes/models.py

@@ -20,7 +20,6 @@ from sklearn.model_selection import train_test_split
 import torch
 from torchvision import models, transforms
 
-from efficientnet.tfkeras import EfficientNetB0
 from data.mnist.mnist_model import Net
 
 def get_xtrain(segs):
@@ -41,9 +40,9 @@ def get_xtrain(segs):
 def process_imagenet_get_model(data):
     """Gets wrapped imagenet model."""
     # Get the vgg16 model, used in the experiments
-    model = EfficientNetB0(weights='imagenet')
-    # model.eval()
-    # model.cuda()
+    model = models.vgg16(pretrained=True)
+    model.eval()
+    model.cuda()
 
     xtest = data['X']
     ytest = data['y'].astype(int)
@@ -73,7 +72,7 @@ def process_imagenet_get_model(data):
                         perturbed_image[segments==i, 1] = background
                         perturbed_image[segments==i, 2] = background
                 perturbed_images.append(transf(perturbed_image)[None, :])
-            perturbed_images = torch.from_numpy(np.concatenate(perturbed_images, axis=0)).float()
+            perturbed_images = torch.from_numpy(np.concatenate(perturbed_images, axis=0)).float().cuda()
             predictions = []
             for q in range(0, perturbed_images.shape[0], batch_size):
                 predictions.append(softmax(model(perturbed_images[q:q+batch_size])).cpu().detach().numpy())

data/mnist/mnist_model.py

@@ -94,7 +94,7 @@ def main():
     parser.add_argument('--save-model', action='store_true', default=False,
                         help='For Saving the current Model')
     args = parser.parse_args()
-    use_cuda = False # not args.no_cuda and torch.cuda.is_available()
+    use_cuda = True # not args.no_cuda and torch.cuda.is_available()
 
     torch.manual_seed(args.seed)