mirror of
https://github.com/modelscope/modelscope.git
synced 2026-07-11 21:09:08 +02:00
fix device error (#419)
This commit is contained in:
@@ -359,4 +359,4 @@ def all_gather(data, group=None):
|
||||
def is_on_same_device(model: torch.nn.Module) -> bool:
|
||||
device_set = set(map(lambda p: p.device.type,
|
||||
model.parameters())) - {'cpu'}
|
||||
return len(device_set) == 1
|
||||
return len(device_set) <= 1
|
||||
|
||||
Reference in New Issue
Block a user