model registry

app.py

@@ -7,6 +7,14 @@ import torch
 import torchvision.models as models
 import torch.nn as nn
 import io
+import wandb
+
+run = wandb.init(project="farmnet", job_type='inference')
+artifact = run.use_artifact("farmnet_model_1:latest", type='model')
+artifact_dir = artifact.download()
+wandb.finish()
+
+
 
 class FarmNet(nn.Module):
     def __init__(self):
@@ -29,7 +37,7 @@ class FarmNet(nn.Module):
         return x
 
 model = FarmNet()  # Replace with your model and architecture
-model.load_state_dict(torch.load('farmnet_model.pth'))
+model.load_state_dict(torch.load('artifacts/farmnet_model_1:v0/farmnet_model.pth'))
 model.eval()  # Set the model to evaluation mode
 
 # Preprocess the image