Saving an attribute

[fabiozappo]

Hi, I am exporting a model using torch neuron but I can't find any reference to save a custom attribute in the model.

For instance I would like to save the dimension of input image as an integer so to be able to get again this value doing something like:

model = torch.jit.load('"/path/to/my/model.pt")
image_size = model.image_size

[jyang-aws]

@fabiozappo as mentioned in torch.jit.load, the pb needs to have saved modules (for a model). a separate tensor is not part of a module. You can get the image_size by either query from the model, say model.layer.weight[0] if it can be inferred from the model. or use something similar to torch.save, which can store dictionary dtype: https://pytorch.org/tutorials/beginner/saving_loading_models.html#save