Avoid non-blocking GPU->CPU copies. (#11288)

Co-authored-by: Jirka Borovec <[email protected]> Co-authored-by: Rohit Gupta <[email protected]> Co-authored-by: Justus Schock <[email protected]> Co-authored-by: Adrian Wälchli <[email protected]>
Lightning-AI · Jan 5, 2022 · 3ab6788 · 3ab6788
1 parent 20fb7cf
commit 3ab6788
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -19,6 +19,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed data fetcher selection ([#11294](https://github.com/PyTorchLightning/pytorch-lightning/pull/11294))
 
 
+- Fixed a race condition that could result in incorrect (zero) values being observed in prediction writer callbacks ([#11288](https://github.com/PyTorchLightning/pytorch-lightning/pull/11288))
+
+
 ## [1.5.7] - 2021-12-21
 
 ### Fixed

diff --git a/pytorch_lightning/utilities/apply_func.py b/pytorch_lightning/utilities/apply_func.py
@@ -34,6 +34,9 @@
     Batch = type(None)
 
 
+_CPU_DEVICES = ("cpu", torch.device("cpu"))
+
+
 def to_dtype_tensor(
     value: Union[int, float, List[Union[int, float]]], dtype: torch.dtype, device: Union[str, torch.device]
 ) -> torch.Tensor:
@@ -268,7 +271,10 @@ def batch_to(data: Any) -> Any:
                 setattr(device_data, field, device_field)
             return device_data
 
-        kwargs = dict(non_blocking=True) if isinstance(data, torch.Tensor) else {}
+        kwargs = {}
+        # Don't issue non-blocking transfers to CPU
+        if isinstance(data, torch.Tensor) and device not in _CPU_DEVICES:
+            kwargs["non_blocking"] = True
         data_output = data.to(device, **kwargs)
         if data_output is not None:
             return data_output