[DeprecationWarning] Addressing warning related to Tensor.pin_memory() argument (pyg-team#9617)

drivanov · web-flow · commit 2c977f597598 · 2024-08-22T04:31:49.000+02:00
This PR resolves the following sets of warnings:
```
test/loader/test_prefetch.py: 10 warnings
  /usr/local/lib/python3.10/dist-packages/torch_geometric/loader/prefetch.py:76: DeprecationWarning: 
The argument 'device' of Tensor.pin_memory() is deprecated. Please do not pass this argument. 
(Triggered internally at /opt/pytorch/pytorch/aten/src/ATen/native/Memory.cpp:46.)
    batch = batch.pin_memory(self.device_helper.device)

test/loader/test_prefetch.py: 10 warnings
  /usr/local/lib/python3.10/dist-packages/torch_geometric/loader/prefetch.py:76: DeprecationWarning: 
The argument 'device' of Tensor.is_pinned() is deprecated. Please do not pass this argument. 
(Triggered internally at /opt/pytorch/pytorch/aten/src/ATen/native/Memory.cpp:31.)
    batch = batch.pin_memory(self.device_helper.device)
```
diff --git a/torch_geometric/loader/prefetch.py b/torch_geometric/loader/prefetch.py
@@ -73,7 +73,7 @@ def non_blocking_transfer(self, batch: Any) -> Any:
         if isinstance(batch, dict):
             return {k: self.non_blocking_transfer(v) for k, v in batch.items()}
 
-        batch = batch.pin_memory(self.device_helper.device)
+        batch = batch.pin_memory()
         return batch.to(self.device_helper.device, non_blocking=True)
 
     def __iter__(self) -> Any:
