fix: templates bug (#456)

* Update model_nn.py * ci issue
microsoft · Oct 28, 2024 · 434a868 · 434a868
1 parent 7b0f043
commit 434a868
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/rdagent/scenarios/kaggle/experiment/digit-recognizer_template/model/model_nn.py b/rdagent/scenarios/kaggle/experiment/digit-recognizer_template/model/model_nn.py
@@ -5,6 +5,8 @@
 import torch.optim as optim
 from torch.utils.data import DataLoader, TensorDataset
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 
 # Define the neural network model with Batch Normalization
 class NeuralNetwork(nn.Module):