(module: torch.nn.Module, args)
| 12 | """ |
| 13 | |
| 14 | def hook(module: torch.nn.Module, args): |
| 15 | if device is not None: |
| 16 | align_device = device |
| 17 | elif len(list(module.parameters())) > 0: |
| 18 | align_device = next(module.parameters()).device |
| 19 | else: |
| 20 | align_device = "cuda" |
| 21 | module.to(align_device) |
| 22 | args = tuple(arg.to(align_device) if isinstance(arg, torch.Tensor) else arg for arg in args) |
| 23 | return args |
| 24 | |
| 25 | def hook_with_kwargs(module: torch.nn.Module, args, kwargs): |
| 26 | if device is not None: |
nothing calls this directly
no outgoing calls
no test coverage detected