def load_state_dict(self, state_dict):
        """
        Load an optimizer state dict.
        Override from PyTorch implementation to avoid casting to FP32.
        """
        if 'loss_scaler' in state_dict:
            # init from the state dict
            if isinstance(state_dict['loss_scaler'], float):
                # new method, restore the float
                self.scaler.loss_scale = state_dict['loss_scaler']
            else:
                # old method, we stored the entire loss scaler
                self.scaler.loss_scale = state_dict['loss_scaler'].loss_scale

        self.optimizer.load_state_dict(state_dict)

        # Hack: PyTorch automatically casts the optimizer state to match the
        # type of the current parameters. But with memory efficient fp16 the
        # params are FP16 while the optimizer state is FP32 and we don't want
        # to cast. A workaround is to manually copy back the original state
        # after the optimizer has been loaded.
        groups = self.optimizer.param_groups
        saved_groups = state_dict['param_groups']
        id_map = {
            old_id: p
            for old_id, p in zip(
                chain(*(g['params'] for g in saved_groups)),
                chain(*(g['params'] for g in groups)),
            )
        }
        for k, v in state_dict['state'].items():
            if k in id_map:
                # make sure when we copy the original state back, we make sure
                # that original state is on the correct device
                param = id_map[k]
                like_device_v = {
                    j: w.to(param.device) if torch.is_tensor(w) else w
                    for j, w in v.items()
                }
                self.optimizer.state[param] = like_device_v

        Load an optimizer state dict.
        Override from PyTorch implementation to avoid casting to FP32.
        
EVA-CLIP/rei/training/fp16.py