Source code for immuneML.ml_methods.util.pytorch_util
from torch import Tensor
from immuneML.data_model.bnp_util import write_yaml
[docs]
def store_weights(model, path):
state_dict = {key: val.tolist() if isinstance(val, Tensor) else val for key, val in model.state_dict().items()}
write_yaml(yaml_dict=state_dict, filename=path)