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]>

Author: 5 people

CHANGELOG.md

@@ -9,11 +9,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
-- 
-
 - Fixed `LightningCLI` race condition while saving the config ([#11199](https://github.com/PyTorchLightning/pytorch-lightning/pull/11199))
-
 - 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

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