retreive trained model from GCS

app.py

@@ -1,10 +1,9 @@
 import torch
 from torch import nn
 from torchvision import models
-import torch.nn.functional as F
 from torchvision.transforms import v2
-
-from PIL import Image
+import os
+import requests
 
 labels = ['Pastel',
             'Yellow Belly',
@@ -53,7 +52,14 @@ def predict(img, confidence):
     # If using GPU
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    checkpoint = torch.load(f'model/model_v8_epoch9.pt', map_location=device)
+    # Download model from GCS
+    model_path = os.environ.get('model_path')
+    response = requests.get(model_path)
+
+    with open('model.pt', 'wb') as f:
+        f.write(response.content)
+
+    checkpoint = torch.load('model.pt', map_location=device)
     densenet.load_state_dict(checkpoint['model_state_dict'])
 
     densenet.eval()