Merge branch 'main' of https://github.com/gpu-mode/discord-cluster-ma…

…nager

.github/workflows/train_workflow.yml

@@ -1,4 +1,4 @@
-name: AMD PyTorch Job
+name: NVIDIA PyTorch Job
 
 on:
   workflow_dispatch:
@@ -10,8 +10,9 @@ on:
 
 jobs:
   train:
-    runs-on: [amdgpu-mi250-x86-64]
-
+    runs-on: [gpumode-nvidia-arc]
+    container:
+      image: nvidia/cuda:12.4.0-devel-ubuntu22.04
     steps:
     - name: Setup Python
       uses: actions/setup-python@v4
@@ -21,7 +22,7 @@ jobs:
     - name: Install dependencies
       run: |
         pip install numpy
-        pip install torch --index-url https://download.pytorch.org/whl/rocm6.2
+        pip install torch
 
     - name: Create training script
       shell: python

train.py

@@ -1,12 +1,8 @@
 import torch
 
-a = torch.Tensor([1, 2, 3, 4, 5]).to("cuda")
-b= torch.Tensor([1, 2, 3, 4, 5]).to("cuda")
-
-if torch.cuda.is_available():
-    gpu_name = torch.cuda.get_device_name(0)
-    print(f"GPU Name: {gpu_name}")
-else:
-    print("No GPU available")
+a = torch.Tensor([1, 2, 3, 4, 5]).cuda()
+b= torch.Tensor([1, 2, 3, 4, 5]).cuda()
 
+print(a)
+print(b)
 print(a + b)