Move model to GPU

CodingTil · Oct 30, 2023 · 04ab49f · 04ab49f
1 parent a0b86c0
commit 04ab49f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/eiuie/fusion_model.py b/eiuie/fusion_model.py
@@ -131,7 +131,7 @@ def __init__(
         self.device = torch.device("cuda" if cuda.is_available() else "cpu")
 
         # Neural Network Model
-        self.net = FusionNet()
+        self.net = FusionNet().to(self.device)
         self.optimizer = optim.Adam(self.net.parameters())
         self.criterion = nn.MSELoss()  # assuming regression task
         self.start_epoch = 0