def save_own_variables(self, store):
        """Saves model's state from `state_dict`.
        `model.parameters` excludes some of model's state like
        `BatchNorm` mean and variance. So, use `state_dict` to obtain
        all of model's state.
        """
        state_dict = self.module.state_dict()
        for key in state_dict.keys():
            store[key] = convert_to_numpy(state_dict[key])
Saves model's state from `state_dict`.
        `model.parameters` excludes some of model's state like
        `BatchNorm` mean and variance. So, use `state_dict` to obtain
        all of model's state.
        
keras/src/utils/torch_utils.py