def forward(self, transition_batch: Transition):
    observation_batch = transition_batch.observation
    action_batch = transition_batch.action
    assert (action_batch.size()[0] == observation_batch.size()[0]), f'Observation batch and action batch should have the same length. Action batch size = {action_batch.size()[0]}, observation batch size = {observation_batch.size()[0]}. Action tensor size = {action_batch.size()[0]}. Observation tensor size = {observation_batch.size()[0]}'
    assert (observation_batch.size() == (observation_batch.size()[0], self._observation_channels_num, self._observation_height, self._observation_width)), f'Observation size should be (Any, {self._observation_channels_num}, {self._observation_height}, {self._observation_width}), instead it is  {observation_batch.size()}'
    assert (action_batch.size()[1] == self._action_size), f'Each action should have size {self._action_size}, not {action_batch.size()[1]}. Tensor has size {action_batch.size()}'
    state_s_0_batch = self.encode_static(observation_batch)
    state_d_0_batch = self.encode_dynamic(observation_batch)
    (state_d_1_batch, reward_d_1_batch) = self.predict_dynamics(state_d_0_batch, action_batch)
    observation_1_batch = self.decode(state_s_0_batch, state_d_1_batch)
    return (observation_1_batch, reward_d_1_batch)
def forward(self, transition_batch: Transition):
    observation_batch = transition_batch.observation
    action_batch = transition_batch.action
    assert (action_batch.size()[0] == observation_batch.size()[0]), f'Observation batch and action batch should have the same length. Action batch size = {action_batch.size()[0]}, observation batch size = {observation_batch.size()[0]}. Action tensor size = {action_batch.size()[0]}. Observation tensor size = {observation_batch.size()[0]}'
    assert (observation_batch.size() == (observation_batch.size()[0], self._observation_channels_num, self._observation_height, self._observation_width)), f'Observation size should be (Any, {self._observation_channels_num}, {self._observation_height}, {self._observation_width}), instead it is  {observation_batch.size()}'
    assert (action_batch.size()[1] == self._action_size), f'Each action should have size {self._action_size}, not {action_batch.size()[1]}. Tensor has size {action_batch.size()}'
    state_s_0_batch = self.encode_static(observation_batch)
    state_d_0_batch = self.encode_dynamic(observation_batch)
    (state_d_1_batch, reward_d_1_batch) = self.predict_dynamics(state_d_0_batch, action_batch)
    observation_1_batch = self.decode(state_s_0_batch, state_d_1_batch)
    return (observation_1_batch, reward_d_1_batch)