Test torch import time

gpu-mode · Nov 5, 2024 · 57a71d0 · 57a71d0
1 parent c1e2b1a
commit 57a71d0
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 8 deletions.
diff --git a/.github/workflows/train_workflow.yml b/.github/workflows/train_workflow.yml
@@ -13,8 +13,7 @@ jobs:
     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: |

diff --git a/train.py b/train.py
@@ -1,8 +1,4 @@
-import numpy
-
-a = numpy.array([1, 2, 3])
-b = numpy.array([4, 5, 6])
-
-c = a + b
+import torch 
 
+a = torch.randn(1)
 print(c)