Enhancement/select device

CHANGELOG.md

@@ -15,6 +15,7 @@
 - [#470](https://github.com/helmholtz-analytics/heat/pull/470) Enhancement: Accelerate distance calculations in kmeans clustering by introduction of new module spatial.distance
 - [#478](https://github.com/helmholtz-analytics/heat/pull/478) `ht.array` now typecasts the local torch tensors if the torch tensors given are not the torch version of the specified dtype + unit test updates
 - [#479](https://github.com/helmholtz-analytics/heat/pull/479) Completion of spatial.distance module to support 2D input arrays of different splittings (None or 0) and different datatypes, also if second input argument is None
+- [#488](https://github.com/helmholtz-analytics/heat/pull/488) Enhancement: test device selection
 
 # v0.2.2
 

heat/cluster/tests/test_kmeans.py

@@ -3,23 +3,33 @@
 
 import heat as ht
 
-if os.environ.get("DEVICE") == "gpu" and ht.torch.cuda.is_available():
+envar = os.getenv("HEAT_USE_DEVICE", "cpu")
+
+if envar == "cpu":
+    ht.use_device("cpu")
+    torch_device = ht.cpu.torch_device
+    heat_device = None
+elif envar == "gpu" and ht.torch.cuda.is_available():
+    ht.use_device("gpu")
+    ht.torch.cuda.set_device(ht.torch.device(ht.gpu.torch_device))
+    torch_device = ht.gpu.torch_device
+    heat_device = None
+elif envar == "lcpu" and ht.torch.cuda.is_available():
     ht.use_device("gpu")
-    ht.torch.cuda.set_device(ht.torch.device(ht.get_device().torch_device))
-else:
+    ht.torch.cuda.set_device(ht.torch.device(ht.gpu.torch_device))
+    torch_device = ht.cpu.torch_device
+    heat_device = ht.cpu
+elif envar == "lgpu" and ht.torch.cuda.is_available():
     ht.use_device("cpu")
-device = ht.get_device().torch_device
-ht_device = None
-if os.environ.get("DEVICE") == "lgpu" and ht.torch.cuda.is_available():
-    device = ht.gpu.torch_device
-    ht_device = ht.gpu
-    ht.torch.cuda.set_device(device)
+    ht.torch.cuda.set_device(ht.torch.device(ht.gpu.torch_device))
+    torch_device = ht.gpu.torch_device
+    heat_device = ht.gpu
 
 
 class TestKMeans(unittest.TestCase):
     def test_fit_iris(self):
         # get some test data
-        iris = ht.load("heat/datasets/data/iris.csv", sep=";")
+        iris = ht.load("heat/datasets/data/iris.csv", sep=";", device=heat_device)
 
         # fit the clusters
         k = 3
@@ -37,7 +47,7 @@ def test_fit_iris(self):
         self.assertIsInstance(kmeans.cluster_centers_, ht.DNDarray)
         self.assertEqual(kmeans.cluster_centers_.shape, (k, iris.shape[1]))
 
-        iris_split = ht.load("heat/datasets/data/iris.csv", sep=";", split=1)
+        iris_split = ht.load("heat/datasets/data/iris.csv", sep=";", split=1, device=heat_device)
         kmeans = ht.cluster.KMeans(n_clusters=k)
 
         with self.assertRaises(NotImplementedError):

heat/core/operations.py

@@ -52,7 +52,7 @@ def __binary_op(operation, t1, t2):
             output_device = None
             output_comm = MPI_WORLD
         elif isinstance(t2, dndarray.DNDarray):
-            t1.gpu() if t2.device.device_type == "gpu" else t1.cpu()
+            t1 = t1.gpu() if t2.device.device_type == "gpu" else t1.cpu()
 
             output_shape = t2.shape
             output_split = t2.split

heat/core/tests/deviceselection.py

@@ -0,0 +1,26 @@
+import torch
+import unittest
+import os
+import heat as ht
+
+envar = os.getenv("HEAT_USE_DEVICE", "cpu")
+
+if envar == "cpu":
+    ht.use_device("cpu")
+    torch_device = ht.cpu.torch_device
+    heat_device = None
+elif envar == "gpu" and torch.cuda.is_available():
+    ht.use_device("gpu")
+    torch.cuda.set_device(torch.device(ht.gpu.torch_device))
+    torch_device = ht.gpu.torch_device
+    heat_device = None
+elif envar == "lcpu" and torch.cuda.is_available():
+    ht.use_device("gpu")
+    torch.cuda.set_device(torch.device(ht.gpu.torch_device))
+    torch_device = ht.cpu.torch_device
+    heat_device = ht.cpu
+elif envar == "lgpu" and torch.cuda.is_available():
+    ht.use_device("cpu")
+    torch.cuda.set_device(torch.device(ht.gpu.torch_device))
+    torch_device = ht.gpu.torch_device
+    heat_device = ht.gpu