From 1fe8bdfd6a62e8c5d20dfa2d41097723b06c128e Mon Sep 17 00:00:00 2001
From: Haoran Pan <167847254+TPLin22@users.noreply.github.com>
Date: Mon, 28 Oct 2024 14:22:04 +0800
Subject: [PATCH] fix: templates bug (#456)

* Update model_nn.py

* ci issue
---
 .../experiment/digit-recognizer_template/model/model_nn.py      | 2 ++
 1 file changed, 2 insertions(+)

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
index 4de18357..d1233f9d 100644
--- 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.nn.functional as F
 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):