ClearML/clearml
1
0
Fork 0
mirror of https://github.com/clearml/clearml synced 2025-02-07 21:33:25 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix issue #292 pytorch-lighting multi node store

Browse Source
This commit is contained in:
allegroai 2021-01-24 09:18:18 +02:00
parent 828af76ffe
commit 381ae21cef
1 changed files with 9 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
clearml/binding/frameworks/pytorch_bind.py
View File

@ -104,10 +104,19 @@ class PatchPyTorchModelIO(PatchBaseModelIO):
@staticmethod
def _save(original_fn, obj, f, *args, **kwargs):
ret = original_fn(obj, f, *args, **kwargs)
# if there is no main task or this is a nested call
if not PatchPyTorchModelIO.__main_task:
return ret
# pytorch-lightning check if rank is zero
if hasattr(obj, 'is_global_zero'):
if not obj.is_global_zero:
return ret
elif hasattr(obj, 'trainer') and hasattr(obj.trainer, 'is_global_zero'):
if not obj.trainer.is_global_zero:
return ret
# noinspection PyBroadException
try:
if isinstance(f, six.string_types):