Calculate hidden state dictionary

def state_dict(self) -> dict:
        """ Calculate hidden state dictionary """
        hidden_state = {}

        if self.optimizer is not None:
            hidden_state['optimizer'] = self.optimizer.state_dict()

        for callback in self.callbacks:
            callback.write_state_dict(self.training_info, hidden_state)

        return hidden_state