fix: templates bug (#456)

* Update model_nn.py * ci issue
2026-07-28 07:57:44 +00:00 · 2024-10-28 14:22:04 +08:00
parent 19317d6d36
commit 1fe8bdfd6a
1 changed files with 2 additions and 0 deletions
@@ -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):