Device agnostic testing
diff --git a/tests/test_layers.py b/tests/test_layers.py
index da061870ab..92f6b683d3 100644
--- a/tests/test_layers.py
+++ b/tests/test_layers.py
@@ -3,6 +3,13 @@
 
 from timm.layers import create_act_layer, set_layer_config
 
+import importlib
+import os
+
+torch_backend = os.environ.get('TORCH_BACKEND')
+if torch_backend is not None:
+    importlib.import_module(torch_backend)
+torch_device = os.environ.get('TORCH_DEVICE', 'cpu')
 
 class MLP(nn.Module):
     def __init__(self, act_layer="relu", inplace=True):
@@ -30,6 +37,9 @@ def _run(x, act_layer=''):
         l = (out - 0).pow(2).sum()
         return l
 
+    x = x.to(device=torch_device)
+    m.to(device=torch_device)
+
     out_me = _run(x)
 
     with set_layer_config(scriptable=True):
diff --git a/tests/test_models.py b/tests/test_models.py
index b1b2bf195a..a6411a7856 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -30,6 +30,17 @@
 from timm.layers import Format, get_spatial_dim, get_channel_dim
 from timm.models import get_notrace_modules, get_notrace_functions
 
+import importlib
+import os
+
+torch_backend = os.environ.get('TORCH_BACKEND')
+if torch_backend is not None:
+    importlib.import_module(torch_backend)
+torch_device = os.environ.get('TORCH_DEVICE', 'cpu')
+timeout = os.environ.get('TIMEOUT')
+timeout120 = int(timeout) if timeout else 120
+timeout300 = int(timeout) if timeout else 300
+
 if hasattr(torch._C, '_jit_set_profiling_executor'):
     # legacy executor is too slow to compile large models for unit tests
     # no need for the fusion performance here
@@ -100,7 +111,7 @@ def _get_input_size(model=None, model_name='', target=None):
 
 
 @pytest.mark.base
-@pytest.mark.timeout(120)
+@pytest.mark.timeout(timeout120)
 @pytest.mark.parametrize('model_name', list_models(exclude_filters=EXCLUDE_FILTERS))
 @pytest.mark.parametrize('batch_size', [1])
 def test_model_forward(model_name, batch_size):
@@ -112,6 +123,8 @@ def test_model_forward(model_name, batch_size):
     if max(input_size) > MAX_FWD_SIZE:
         pytest.skip("Fixed input size model > limit.")
     inputs = torch.randn((batch_size, *input_size))
+    inputs = inputs.to(torch_device)
+    model.to(torch_device)
     outputs = model(inputs)
 
     assert outputs.shape[0] == batch_size
@@ -119,7 +132,7 @@ def test_model_forward(model_name, batch_size):
 
 
 @pytest.mark.base
-@pytest.mark.timeout(120)
+@pytest.mark.timeout(timeout120)
 @pytest.mark.parametrize('model_name', list_models(exclude_filters=EXCLUDE_FILTERS, name_matches_cfg=True))
 @pytest.mark.parametrize('batch_size', [2])
 def test_model_backward(model_name, batch_size):
@@ -133,6 +146,8 @@ def test_model_backward(model_name, batch_size):
     model.train()
 
     inputs = torch.randn((batch_size, *input_size))
+    inputs = inputs.to(torch_device)
+    model.to(torch_device)
     outputs = model(inputs)
     if isinstance(outputs, tuple):
         outputs = torch.cat(outputs)
@@ -147,7 +162,7 @@ def test_model_backward(model_name, batch_size):
 
 
 @pytest.mark.cfg
-@pytest.mark.timeout(300)
+@pytest.mark.timeout(timeout300)
 @pytest.mark.parametrize('model_name', list_models(
     exclude_filters=EXCLUDE_FILTERS + NON_STD_FILTERS, include_tags=True))
 @pytest.mark.parametrize('batch_size', [1])
@@ -155,6 +170,7 @@ def test_model_default_cfgs(model_name, batch_size):
     """Run a single forward pass with each model"""
     model = create_model(model_name, pretrained=False)
     model.eval()
+    model.to(torch_device)
     state_dict = model.state_dict()
     cfg = model.default_cfg
 
@@ -169,7 +185,7 @@ def test_model_default_cfgs(model_name, batch_size):
             not any([fnmatch.fnmatch(model_name, x) for x in EXCLUDE_FILTERS]):
         # output sizes only checked if default res <= 448 * 448 to keep resource down
         input_size = tuple([min(x, MAX_FWD_OUT_SIZE) for x in input_size])
-        input_tensor = torch.randn((batch_size, *input_size))
+        input_tensor = torch.randn((batch_size, *input_size), device=torch_device)
 
         # test forward_features (always unpooled)
         outputs = model.forward_features(input_tensor)
@@ -180,12 +196,14 @@ def test_model_default_cfgs(model_name, batch_size):
 
         # test forward after deleting the classifier, output should be poooled, size(-1) == model.num_features
         model.reset_classifier(0)
+        model.to(torch_device)
         outputs = model.forward(input_tensor)
         assert len(outputs.shape) == 2
         assert outputs.shape[1] == model.num_features
 
         # test model forward without pooling and classifier
         model.reset_classifier(0, '')  # reset classifier and set global pooling to pass-through
+        model.to(torch_device)
         outputs = model.forward(input_tensor)
         assert len(outputs.shape) == 4
         if not isinstance(model, (timm.models.MobileNetV3, timm.models.GhostNet, timm.models.RepGhostNet, timm.models.VGG)):
@@ -195,6 +213,7 @@ def test_model_default_cfgs(model_name, batch_size):
         if 'pruned' not in model_name:  # FIXME better pruned model handling
             # test classifier + global pool deletion via __init__
             model = create_model(model_name, pretrained=False, num_classes=0, global_pool='').eval()
+            model.to(torch_device)
             outputs = model.forward(input_tensor)
             assert len(outputs.shape) == 4
             if not isinstance(model, (timm.models.MobileNetV3, timm.models.GhostNet, timm.models.RepGhostNet, timm.models.VGG)):
@@ -218,13 +237,14 @@ def test_model_default_cfgs(model_name, batch_size):
 
 
 @pytest.mark.cfg
-@pytest.mark.timeout(300)
+@pytest.mark.timeout(timeout300)
 @pytest.mark.parametrize('model_name', list_models(filter=NON_STD_FILTERS, exclude_filters=NON_STD_EXCLUDE_FILTERS, include_tags=True))
 @pytest.mark.parametrize('batch_size', [1])
 def test_model_default_cfgs_non_std(model_name, batch_size):
     """Run a single forward pass with each model"""
     model = create_model(model_name, pretrained=False)
     model.eval()
+    model.to(torch_device)
     state_dict = model.state_dict()
     cfg = model.default_cfg
 
@@ -232,7 +252,7 @@ def test_model_default_cfgs_non_std(model_name, batch_size):
     if max(input_size) > 320:  # FIXME const
         pytest.skip("Fixed input size model > limit.")
 
-    input_tensor = torch.randn((batch_size, *input_size))
+    input_tensor = torch.randn((batch_size, *input_size), device=torch_device)
     feat_dim = getattr(model, 'feature_dim', None)
 
     outputs = model.forward_features(input_tensor)
@@ -246,6 +266,7 @@ def test_model_default_cfgs_non_std(model_name, batch_size):
 
     # test forward after deleting the classifier, output should be poooled, size(-1) == model.num_features
     model.reset_classifier(0)
+    model.to(torch_device)
     outputs = model.forward(input_tensor)
     if isinstance(outputs,  (tuple, list)):
         outputs = outputs[0]
@@ -254,6 +275,7 @@ def test_model_default_cfgs_non_std(model_name, batch_size):
     assert outputs.shape[feat_dim] == model.num_features, 'pooled num_features != config'
 
     model = create_model(model_name, pretrained=False, num_classes=0).eval()
+    model.to(torch_device)
     outputs = model.forward(input_tensor)
     if isinstance(outputs, (tuple, list)):
         outputs = outputs[0]
@@ -297,7 +319,7 @@ def test_model_features_pretrained(model_name, batch_size):
 
 
 @pytest.mark.torchscript
-@pytest.mark.timeout(120)
+@pytest.mark.timeout(timeout120)
 @pytest.mark.parametrize(
     'model_name', list_models(exclude_filters=EXCLUDE_FILTERS + EXCLUDE_JIT_FILTERS, name_matches_cfg=True))
 @pytest.mark.parametrize('batch_size', [1])
@@ -312,6 +334,7 @@ def test_model_forward_torchscript(model_name, batch_size):
     model.eval()
 
     model = torch.jit.script(model)
+    model.to(torch_device)
     outputs = model(torch.randn((batch_size, *input_size)))
 
     assert outputs.shape[0] == batch_size
diff --git a/tests/test_optim.py b/tests/test_optim.py
index 9bdfd6825d..38f625fb42 100644
--- a/tests/test_optim.py
+++ b/tests/test_optim.py
@@ -15,6 +15,13 @@
 
 from timm.optim import create_optimizer_v2
 
+import importlib
+import os
+
+torch_backend = os.environ.get('TORCH_BACKEND')
+if torch_backend is not None:
+    importlib.import_module(torch_backend)
+torch_device = os.environ.get('TORCH_DEVICE', 'cuda')
 
 # HACK relying on internal PyTorch test functionality for comparisons that I don't want to write
 torch_tc = TestCase()
@@ -61,7 +68,7 @@ def _test_state_dict(weight, bias, input, constructor):
 
     def fn_base(optimizer, weight, bias):
         optimizer.zero_grad()
-        i = input_cuda if weight.is_cuda else input
+        i = input_device if weight.device.type != 'cpu' else input
         loss = (weight.mv(i) + bias).pow(2).sum()
         loss.backward()
         return loss
@@ -97,28 +104,30 @@ def fn_base(optimizer, weight, bias):
 
     # Check that state dict can be loaded even when we cast parameters
     # to a different type and move to a different device.
-    if not torch.cuda.is_available():
+    if torch_device == 'cpu':
+        return
+    elif torch_device == 'cuda' and not torch.cuda.is_available():
         return
 
     with torch.no_grad():
-        input_cuda = Parameter(input.clone().detach().float().cuda())
-        weight_cuda = Parameter(weight.clone().detach().cuda())
-        bias_cuda = Parameter(bias.clone().detach().cuda())
-    optimizer_cuda = constructor(weight_cuda, bias_cuda)
-    fn_cuda = functools.partial(fn_base, optimizer_cuda, weight_cuda, bias_cuda)
+        input_device = Parameter(input.clone().detach().float().to(torch_device))
+        weight_device = Parameter(weight.clone().detach().to(torch_device))
+        bias_device = Parameter(bias.clone().detach().to(torch_device))
+    optimizer_device = constructor(weight_device, bias_device)
+    fn_device = functools.partial(fn_base, optimizer_device, weight_device, bias_device)
 
     state_dict = deepcopy(optimizer.state_dict())
     state_dict_c = deepcopy(optimizer.state_dict())
-    optimizer_cuda.load_state_dict(state_dict_c)
+    optimizer_device.load_state_dict(state_dict_c)
 
     # Make sure state dict wasn't modified
     torch_tc.assertEqual(state_dict, state_dict_c)
 
     for _i in range(20):
         optimizer.step(fn)
-        optimizer_cuda.step(fn_cuda)
-        torch_tc.assertEqual(weight, weight_cuda)
-        torch_tc.assertEqual(bias, bias_cuda)
+        optimizer_device.step(fn_device)
+        torch_tc.assertEqual(weight, weight_device)
+        torch_tc.assertEqual(bias, bias_device)
 
     # validate deepcopy() copies all public attributes
     def getPublicAttr(obj):
@@ -152,12 +161,15 @@ def _test_basic_cases(constructor, scheduler_constructors=None):
         scheduler_constructors
     )
     # CUDA
-    if not torch.cuda.is_available():
+    if torch_device == 'cpu':
         return
+    elif torch_device == 'cuda' and not torch.cuda.is_available():
+        return
+
     _test_basic_cases_template(
-        torch.randn(10, 5).cuda(),
-        torch.randn(10).cuda(),
-        torch.randn(5).cuda(),
+        torch.randn(10, 5).to(torch_device),
+        torch.randn(10).to(torch_device),
+        torch.randn(5).to(torch_device),
         constructor,
         scheduler_constructors
     )

This PR changes the test suite to allow running the tests on any hardware.

In order to do this I tried to make the test suite device agnostic and allow setting the device and backend through environment variables. This is hopefully a future proof way for allowing the test suite to be run on any hardware.

I've also added the functionality to change timeouts on tests that run on hardware.

In all the changes I've tried to keep previous functionality working if using the CI in the "normal" way. If the environment variables are not used, previous values are used as defaults.

Example usage:
`TORCH_DEVICE=fancy_new_hw TORCH_BACKEND=new_torch_backend pytest tests/`

