feat: Updated src/main.py

sweepai · Oct 24, 2023 · c80fd1e · c80fd1e
1 parent d1da9cd
commit c80fd1e
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/main.py b/src/main.py
@@ -1,3 +1,9 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+
 class MNISTTrainer:
  def __init__(self):
  self.transform = transforms.Compose([