Skip to content

Commit a914e4a

Browse files
committed
Don't use a non-blocking copy.
It's fragile, harder to manage lifetimes, and doesn't win us measurable gains.
1 parent 672e0f0 commit a914e4a

1 file changed

Lines changed: 4 additions & 4 deletions

File tree

src/mss/screenshot.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -338,10 +338,10 @@ def to_torch( # noqa: PLR0912
338338
rv = torch.frombuffer(self._raw, dtype=torch.uint8)
339339
rv = rv.reshape((self.height, self.width, 4))
340340

341-
# Move the data to the desired device. If no copy is needed, this is a no-op. PyTorch using CUDA can do this
342-
# transfer without blocking; the other devices can't. (Well, they technically can, but then our subsequent ops
343-
# may corrupt data unless we synchronize explicitly.)
344-
rv = rv.to(device=torch_device, non_blocking=(torch_device.type == "cuda"))
341+
# Move the data to the desired device. If no copy is needed, this is a no-op.
342+
# We don't use a non-blocking copy because it can be fragile, hard to manage lifetimes, and doesn't win us
343+
# measurable gains.
344+
rv = rv.to(device=torch_device)
345345

346346
if channels == "BGRA":
347347
pass

0 commit comments

Comments
 (0)