do not use assing=True for nn.LayerNorm

This commit is contained in:
Won-Kyu Park 2024-09-25 22:12:46 +09:00
parent 5f3314ec43
commit 4ad5f22c7b
No known key found for this signature in database
GPG Key ID: 53AA79C8C9535D15

View File

@ -176,7 +176,7 @@ class LoadStateDictOnMeta(ReplaceHelper):
def load_from_state_dict(original, module, state_dict, prefix, *args, **kwargs):
used_param_keys = []
if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d, torch.nn.GroupNorm, torch.nn.LayerNorm)):
if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d, torch.nn.GroupNorm,)):
# HACK add assign=True to local_metadata for some cases
args[0]['assign_to_params_buffers'] = True