(obj: flow.nn.Module, dest: MAP_LOCATION)
| 250 | |
| 251 | |
| 252 | def module_to(obj: flow.nn.Module, dest: MAP_LOCATION) -> "oneflow.nn.Module": |
| 253 | restore_location = _get_restore_location(dest) |
| 254 | # for nn.Module object, we will use a tensor to get the device |
| 255 | # to support dest with a Callable type |
| 256 | device = restore_location(flow.tensor([0])).device |
| 257 | obj.to(device) |
| 258 | return obj |
| 259 | |
| 260 | |
| 261 | def _map_location(obj: Any, map_location: MAP_LOCATION): |
no test coverage detected