From 2f7abaaada8b91d3bbfc2f4ca52e6efb4ad1cdc6 Mon Sep 17 00:00:00 2001
From: Tamino Huxohl <thuxohl@techfak.uni-bielefeld.de>
Date: Fri, 16 Dec 2022 16:02:55 +0100
Subject: [PATCH] distance training moves model to correct device
---
mu_map/training/distance.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mu_map/training/distance.py b/mu_map/training/distance.py
index e607dff..5c44605 100644
--- a/mu_map/training/distance.py
+++ b/mu_map/training/distance.py
@@ -228,7 +228,7 @@ if __name__ == "__main__":
logger=logger,
)
- model = UNet(in_channels=1, features=args.features)
+ model = UNet(in_channels=1, features=args.features).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, betas=(0.5, 0.999))
lr_scheduler = (
torch.optim.lr_scheduler.StepLR(
--
GitLab