I am in a situation were I initialize and train a model using some datasets. I then save the model state dict. I would like to be able to load the model without having to load the data again. However because of the dimensions mismatch when I call
load_state_dict I don’t seem to be able to do that.
class Model(nn.Module): def __init__( self, layer_list, x_mean: Optional[np.array] = None, x_std: Optional[np.array] = None, **_ ): super(Model, self).__init__() self.register_buffer('x_mean', None if x_mean is None else torch.FloatTensor(x_mean)) self.register_buffer('x_std', None if x_std is None else torch.FloatTensor(x_std)) self.sequential = get_sequential(layer_list)
After training the
x_std, … are filled with tensor of certain shapes. However when I load the model, I ‘don’t know’ the shape of those, and I would like to avoid to search in the data what are the shape of x_mean and x_std.
In this case I would like to load that model, just using the
layer_list and the function