(device: torch.device)
| 32 | |
| 33 | |
| 34 | def synchronize_device(device: torch.device): |
| 35 | if device.type == "cuda": |
| 36 | torch.cuda.synchronize() |
| 37 | elif device.type == "xpu": |
| 38 | torch.xpu.synchronize() |
| 39 | elif device.type == "mps": |
| 40 | torch.mps.synchronize() |
| 41 | |
| 42 | |
| 43 | def swap_weight_devices_cuda(device: torch.device, layer_to_cpu: nn.Module, layer_to_cuda: nn.Module): |
no outgoing calls
no test coverage detected