update predict_image.py

predict_image.py

@@ -27,14 +27,17 @@ def load_model(model_name):
     if model_name == 'effb0':
         model = EffnetModel() 
         fname = download_weights(effb0_net_url)
+        print('loaded effnet')
     elif model_name == 'res18':
         model = ResnetModel
         fname = download_weights(res18_net_url)
+        print('loaded resnet')
     else:
         raise ValueError('Enter correct model_name')
 
     # loading pretrained model
     state_dict = torch.load(fname)
+    print(type(state_dict))
     model.load_state_dict(state_dict['weights'])
     return model