def METHOD_NAME(self, state_dict, name):
    prefix = name + "." if name != "" else ""
    for k, v in state_dict.items():
        if k.endswith(prefix + "weight"):
            if v.dim() == 3 and v.size(1) == 1:
                state_dict[k] = v.squeeze(1)