ci: cleanup torch backend tests (#16750)

.github/workflows/test.yml

@@ -77,45 +77,14 @@ jobs:
       run: |
         sudo apt update || true
         sudo apt install -y --no-install-recommends ninja-build
-    - name: Test one op
-      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_tiny.py TestTiny.test_plus
     - name: Test ResNet-18
       run: DEBUG=2 python3 extra/torch_backend/example.py
-    - name: custom tests
-      run: python3 -m pytest -n auto extra/torch_backend/test.py --durations=20
     - name: Test one op in torch tests
       run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
     - name: Test Ops with TINY_BACKEND
       run: DEV=CPU:LLVM LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/backend/test_ops.py --durations=20
-    - name: Test in-place operations on views
-      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
-    - name: Test multi-gpu
-      run: DEV=CPU:LLVM GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
-    - name: Test kernel fusion
-      run: python3 extra/torch_backend/test_kernel_fusion.py
-
-
-  torchbackendmore:
-    name: Torch Backend Tests More
-    runs-on: *linux
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_unit
-        llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
-    - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: STEPS=20 DEV=CPU TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
-    - name: Test some torch tests (expect failure)
-      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
+    - name: Custom tests
+      run: DEV=CPU:LLVM GPUS=4 TINY_BACKEND=1 python3 -m pytest -nauto extra/torch_backend/test.py extra/torch_backend/test_inplace.py extra/torch_backend/test_multigpu.py extra/torch_backend/test_kernel_fusion.py --durations=20
 
   bepython:
     name: Python Backend

extra/torch_backend/test_inplace.py

@@ -1,10 +1,14 @@
 import unittest
 import torch
 import tinygrad.nn.torch
-torch.set_default_device("tiny")
 import numpy as np
 
 class TestTorchBackendInplace(unittest.TestCase):
+  def setUp(self):
+    self._prev_device = torch.get_default_device()
+    torch.set_default_device("tiny")
+  def tearDown(self): torch.set_default_device(self._prev_device)
+
   def test_zero(self):
     a = torch.ones(4)
     a.zero_()

extra/torch_backend/test_multigpu.py

@@ -2,11 +2,15 @@ import unittest
 from tinygrad.helpers import getenv
 import torch
 import tinygrad.nn.torch
-torch.set_default_device("tiny")
 import numpy as np
 
 @unittest.skipIf(getenv("GPUS",1)<=1, "only single GPU")
 class TestTorchBackendMultiGPU(unittest.TestCase):
+  def setUp(self):
+    self._prev_device = torch.get_default_device()
+    torch.set_default_device("tiny")
+  def tearDown(self): torch.set_default_device(self._prev_device)
+
   def test_transfer(self):
     a = torch.Tensor([[1,2],[3,4]]).to("tiny:0")
     b = torch.Tensor([[3,2],[1,0]]).to("tiny:1")