How to extract the attributes of a saved model before loading the whole model

ZSheikhb: <p>I have the following class. I’d like to save this class in one script and load it in another one. However, I don’t know how I can correctly initialize the attributes of a class when I load it somewhere else?!</p>
<pre><code class="lang-auto">class Dataset(nn.Module):

  def __init__(
      self,
      observation_spec,
      action_spec,
      size,
      ):
    super(Dataset, self).__init__()
    self._size = size
    obs_shape = list(observation_spec.shape)
    obs_type = observation_spec.dtype
    action_shape = list(action_spec.shape)
    action_type = action_spec.dtype
    self._s1 = self._zeros([size] + obs_shape, obs_type)
    self._s2 = self._zeros([size] + obs_shape, obs_type)
    self._a1 = self._zeros([size] + action_shape, action_type)
    self._a2 = self._zeros([size] + action_shape, action_type)
    self._discount = self._zeros([size], torch.float32)
    self._reward = self._zeros([size], torch.float32)
    self._data = Transition(
        s1=self._s1, s2=self._s2, a1=self._a1, a2=self._a2,
        discount=self._discount, reward=self._reward)
    self._current_size = torch.autograd.Variable(torch.tensor(0), requires_grad=False)
    self._current_idx = torch.autograd.Variable(torch.tensor(0), requires_grad=False)
    self._capacity = torch.autograd.Variable(torch.tensor(self._size))
    self._config = collections.OrderedDict(
        observation_spec=observation_spec,
        action_spec=action_spec,
        size=size

  @property
  def config(self):
    return self._config

  @property
  def data(self):
    return self._data

  @property
  def capacity(self):
    return self._size

  @property
  def size(self):
    return self._current_size.numpy()

  def _zeros(self, shape, dtype):
    """Create a variable initialized with zeros."""
    return torch.autograd.Variable(torch.zeros(shape, dtype = dtype))
#save the model/class
  assert data.size == data.capacity
  data_ckpt_name = os.path.join(log_dir, 'data_{}.pt'.format(env_name))
  torch.save([data.capacity, data.state_dict()], data_ckpt_name)
  whole_data_ckpt_name = os.path.join(log_dir, 'data_{}.pth'.format(env_name))
  with open( whole_data_ckpt_name, 'wb') as filehandler: 
    pickle.dump(data, filehandler)
</code></pre>
<p>when I tried to load this class based on <a href="https://stackoverflow.com/questions/69903636/how-can-i-load-a-model-in-pytorch-without-having-to-remember-the-parameters-used" rel="noopener nofollow ugc">this answer</a> with its attribute inside another script</p>
<pre><code class="lang-auto">  dm_env = gym.spec(env_name).make()
  env = alf_gym_wrapper.AlfGymWrapper(dm_env)
  observation_spec = env.observation_spec()
  action_spec = env.action_spec()

  # Prepare data.
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  data_ckpt_name = os.path.join(data_file, 'data_{}.pt'.format(env_name))
  whole_data_ckpt_name = os.path.join(data_file, 'data_{}.pth'.format(env_name))
  
  data_size, state = torch.load(data_ckpt_name, map_location=device)
  full_data = dc.Dataset(observation_spec, action_spec, data_size)
  with open(whole_data_ckpt_name, 'wb') as filehandler:
      full_data = pickle.load(filehandler)
  full_data.load_state_dict( state)
  print(f"loaded data : {full_data.size}")   
</code></pre>
<p>I got this error message</p>
<pre><code class="lang-auto">  raise proxy.with_traceback(exception.__traceback__) from None
  File "/home/lib/python3.8/site-packages/gin/config.py", line 1582, in gin_wrapper
    return fn(*new_args, **new_kwargs)
  File "train_eval_offline.py", line 79, in train_eval_offline
    full_data = pickle.load(filehandler)
io.UnsupportedOperation: read
  In call to configurable 'train_eval_offline' (&lt;function train_eval_offline at 0x2b5ceb9dc8b0&gt;)
</code></pre>
<p>I am wondering how I can extract the <code>size</code> attribute of the class from the saved model when I load it, in order to properly initialize this attribute?</p>