Upload folder using huggingface_hub

device.py

@@ -0,0 +1,7 @@
+import torch
+
+device: str
+
+if torch.cuda.is_available(): device = 'cuda'
+elif torch.mps.is_available(): device = 'mps'
+else: device = 'cpu'

model.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:75c98e3b3f34bf836379b5c428cdba0f48794be3ea418b7b9418b9754ec50231
+oid sha256:287330a63e005470ced2f2678a429438d4194f83c524e22b8c986c39c800a257
 size 172209

run.py

@@ -1,12 +1,13 @@
 import torch
+from device import device
 
-model = torch.load("model.pth", weights_only=False)
+model = torch.load("model.pth", weights_only=False).to(device)
 
 def run(test):
     with torch.no_grad():
-        test_data = torch.tensor([test], dtype=torch.float)
+        test_data = torch.tensor([test], dtype=torch.float).to(device)
         predictions = model(test_data)
-        return predictions.squeeze().numpy()
+        return predictions.squeeze().cpu().numpy()
 
 if __name__ == '__main__':
     x, y = map(int, input().split())

train.py

@@ -3,6 +3,7 @@ import torch.nn as nn
 import torch.optim as optim
 import numpy as np
 from MultiplicationNet import MultiplicationNet
+from device import device
 
 def generate_data(num_samples, min_val=0, max_val=100):
     x1 = np.random.randint(min_val, max_val, size=(num_samples, 1))
@@ -16,11 +17,11 @@ def train():
     learning_rate = 0.01
 
     x, y = generate_data(num_samples)
-    x_train = torch.tensor(x, dtype=torch.float)
-    y_train = torch.tensor(y, dtype=torch.float)
+    x_train = torch.tensor(x, dtype=torch.float).to(device)
+    y_train = torch.tensor(y, dtype=torch.float).to(device)
 
-    model = MultiplicationNet()
-    criterion = nn.MSELoss()
+    model = MultiplicationNet().to(device)
+    criterion = nn.MSELoss().to(device)
     optimizer = optim.Adam(model.parameters(), lr=learning_rate)
     scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.95)