[Feature] Extended support to include NPU use cases in addition to existing CUDA-compatible scenarios (#1460)

Co-authored-by: chenhao388 <[email protected]>
Co-authored-by: Vincent Moens <[email protected]>

Author: 3 people

test/_utils_internal.py

@@ -6,6 +6,7 @@
 import pathlib
 import shutil
 import tempfile
+import warnings
 
 import numpy as np
 import torch
@@ -32,6 +33,10 @@ def prod(sequence):
         return int(np.prod(sequence))
 
 
+def is_npu_available():
+    return hasattr(torch, "npu") and torch.npu.is_available()
+
+
 def get_available_devices():
     devices = [torch.device("cpu")]
     n_cuda = torch.cuda.device_count()
@@ -40,6 +45,16 @@ def get_available_devices():
             devices += [torch.device(f"cuda:{i}")]
             if i == 1:
                 break
+    if is_npu_available():
+        warnings.warn(
+            "torch_npu is an experimental feature and not currently included in tensordict CI/CD."
+        )
+        n_npu = torch.npu.device_count()
+        if n_npu > 0:
+            for i in range(n_npu):
+                devices += [torch.device(f"npu:{i}")]
+                if i == 1:
+                    break
     # TODO: MPS and NPU would be worth considering but it's a lot of work
     # for example, many ops are tested with various dtypes but MPS struggles with
     # float64. Shared mem can also cause trouble.