Add torch and GPU

.github/workflows/train_workflow.yml

@@ -9,12 +9,11 @@ on:
 
 jobs:
   train:
-    runs-on: ubuntu-latest
+    runs-on: [gpumode-nvidia-arc, amd-mi250] # Add more GPUs as we get them here
     steps:
       - name: Install dependencies
         run: |
-          pip install numpy
-          # pip install torch - need to find a way to cache this otherwise it will take a long time to install
+          pip install numpy torch
 
       - name: Create and run training script
         run: |

train.py

@@ -1,8 +1,12 @@
-import numpy
+import torch
 
-a = numpy.array([1, 2, 3])
-b = numpy.array([4, 5, 6])
 
+
+# Vector addition
+a = torch.tensor([1, 2, 3]).cuda()
+b = torch.tensor([4, 5, 6]).cuda()
 c = a + b
 
-print(c)
+print(a)
+print(b)
+print(c)