Hi @<1543766544847212544:profile|SorePelican79> ! You could use the following workaround:
from clearml import Task
from clearml.binding.frameworks import WeightsFileHandler
import torch
def filter_callback(
callback_type: WeightsFileHandler.CallbackType,
model_info: WeightsFileHandler.ModelInfo,
):
print(model_info.__dict__)
if (
callback_type == WeightsFileHandler.CallbackType.save
and "filter_out.pt" in model_info.local_model_path
):
return None
return model_info
if __name__ == "__main__":
Task.init(project_name="example", task_name="filter out")
WeightsFileHandler.add_pre_callback(filter_callback)
filter_out_model = torch.nn.Module()
dont_filter_out_model = torch.nn.Module()
torch.save(filter_out_model, "filter_out.pt")
torch.save(dont_filter_out_model, "dont_filter_out_model.pt")
As you can see, you can use WeightsFileHandler
after the task has been initialized