    def test_state_dict_data_parallel_fix(self):
        '''
        test the state_dict_data_parallel_fix function
        Returns: None
        '''
        x = {'module.text_model.embeddings.word_embeddings.weight': torch.randn(2, 4),
             'module.text_model.embeddings.position_embeddings.weight': torch.randn(2, 4)}
        y = {'text_model.embeddings.word_embeddings.weight': torch.randn(2, 4),
             'text_model.embeddings.position_embeddings.weight': torch.randn(2, 4)}
        out = state_dict_data_parallel_fix(x, y)
        self.assertEqual(y.keys(), out.keys())
test the state_dict_data_parallel_fix function
Returns: None
tests/unittests/models/frozen_in_time/test_frozen_utils.py