set device for model

app.py

@@ -18,7 +18,7 @@ def request_pydantic_schema(_: type, handler: GetCoreSchemaHandler):
 Request.__get_pydantic_core_schema__ = request_pydantic_schema
 
 device = set_device()
-model = load_model()
+model = load_model(device=device)
 model.eval()
 
 transform = T.Compose([

app/main.py

@@ -22,7 +22,7 @@ app = FastAPI()
 
 device = set_device()
 log.info("Using device: %s", device)
-model = load_model()
+model = load_model(device=device)
 log.info("Model loaded. App is ready.")
 
 model.eval()

src/model_config.py

@@ -1,4 +1,6 @@
-from train import Autoencoder, DEVICE
+# src/model_config.py
+
+from train import Autoencoder
 import torch
 from huggingface_hub import hf_hub_download
 import logging
@@ -12,15 +14,15 @@ def set_device():
     return device
 
 # Load model
-def load_model(h: int = 32):
+def load_model(device, h: int = 32):
     # model_path = "models/autoencoder_mnist.pth" # Use local model
     model_path = hf_hub_download(
         repo_id="dmtschulz/anomaly-detection-model",
         filename="autoencoder_mnist.pth"
     )
 
-    model = Autoencoder(h).to(DEVICE)
-    state_dict = torch.load(model_path, map_location=DEVICE)
+    model = Autoencoder(h).to(device)
+    state_dict = torch.load(model_path, map_location=device)
     model.load_state_dict(state_dict)
     model.eval()
     log.info("Model loaded from %s", model_path)