changed weight loading

2026-04-29 14:53:47 +02:00
parent c31abed2f0
commit a9cb14cf42
2 changed files with 1 additions and 1 deletions
@@ -38,7 +38,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

 model = Bird_CNN(c_in=3, c_hidden=15, c_out=7)
 full_path = os.path.join(BUILD_PATH, "bird_cnn")
-model.load_state_dict(torch.load(full_path, weights_only=False))
+model.load_state_dict(torch.load(full_path, map_location=torch.device(device)))
 model.to(device)
 model.eval()