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):