fix: todevice ignores non_blocking for nested containers

[andrewwhitecdw]

Problem

fix: todevice ignores non_blocking for nested containers

Fix

Replace:

    if isinstance(batch, dict):
        return {k: todevice(v, device) for k, v in batch.items()}

    if isinstance(batch, (tuple, list)):
        return type(batch)(todevice(x, device) for x in batch)

with:

    if isinstance(batch, dict):
        return {k: todevice(v, device, callback=callback, non_blocking=non_blocking) for k, v in batch.items()}

    if isinstance(batch, (tuple, list)):
        return type(batch)(todevice(x, device, callback=callback, non_blocking=non_blocking) for x in batch)

Files changed

• dust3r/utils/device.py