Single optimization step for a model

def optimizer_step(self, batch_info, device, model, rollout):
        """ Single optimization step for a model """
        batch_info.optimizer.zero_grad()

        batch_result = self.calculate_gradient(batch_info=batch_info, device=device, model=model, rollout=rollout)

        clip_gradients(batch_result, model, self.max_grad_norm)

        batch_info.optimizer.step(closure=None)

        self.post_optimization_step(batch_info, device, model, rollout)

        return batch_result