Bug fix.

src/lambo/learning/train.py

@@ -223,7 +223,7 @@ def pretrain(model, train_dataloader, test_dataloader, epochs, device='cpu'):
     optimizer = Adam(model.parameters(), lr=learning_rate)
     
     print("Pretraining")
-    test_loop_pretraining(test_dataloader, model)
+    test_loop_pretraining(test_dataloader, model, device)
     for t in range(epochs):
         print(f"Epoch {t + 1}\n-------------------------------")
         train_loop(train_dataloader, model, optimizer, device)