def forward(self, obs):
        '''
            Forward pass of actor network. Input assume to be output of CNN
            and/or LSTM feature extractor
            Args:
                obs: batched tensor denotes the states
        '''
        obs_shape = obs.size()
        if_high_dim = (len(obs_shape) == 3)
        if if_high_dim: 
            obs = obs.view(-1, obs_shape[2])

        v = self.model(obs)

        if if_high_dim:
            v = v.view(obs_shape[0], obs_shape[1], 1)
        return v

            Forward pass of actor network. Input assume to be output of CNN
            and/or LSTM feature extractor
            Args:
                obs: batched tensor denotes the states
        
surreal/model/model_builders/builders.py