Revert "Set device for torch tensors with gpu > 1 (#132)" (#134)

This reverts commit 8782c9d.
NVIDIA-Merlin · Apr 14, 2023 · 014b658 · 014b658
1 parent 8782c9d
commit 014b658
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 8 deletions.
diff --git a/merlin/dataloader/loader_base.py b/merlin/dataloader/loader_base.py
@@ -447,8 +447,7 @@ def _to_tensor(self, df_or_series):
         if self.device == "cpu":
             tensor = df_or_series.to_numpy()
         else:
-            with cupy.cuda.Device(self.device):
-                tensor = df_or_series.to_cupy()
+            tensor = df_or_series.to_cupy()
 
         return tensor
 

diff --git a/merlin/dataloader/torch.py b/merlin/dataloader/torch.py
@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import contextlib
 from functools import partial
 
 from merlin.core.compat.torch import torch as th
@@ -119,11 +118,6 @@ def map(self, fn):
 
         return self
 
-    def _get_device_ctx(self, dev):
-        if dev == "cpu" or not th:
-            return contextlib.nullcontext()
-        return th.cuda.device(f"cuda:{dev}")
-
 
 class DLDataLoader(th.utils.data.DataLoader):
     """