changed to efficientnet_b1

bayes/models.py

@@ -40,7 +40,7 @@ def get_xtrain(segs):
 def process_imagenet_get_model(data):
     """Gets wrapped imagenet model."""
     # Get the vgg16 model, used in the experiments
-    model = models.efficientnet_b5(pretrained=True)
+    model = models.efficientnet_b1(pretrained=True)
     model.eval()
     # model.cuda()