def convnext_timm_tiny_2d_to_3d(model, state_dict, ignore_head=True, normalize=True):
    """
    inputs:
        model: nn.Module, the from 'scratch' model
        state_dict: dict, from the pretrained weights
        ignore_head: bool, whether to inflate weights in the head (or keep scratch weights).
            If number of classes changes (eg, imagenet to hmdb51), then you need to use this.
        normalize: bool, if set to True (default), it inflates with a factor of 1, and if
             set to False it inflates with a factor of 1/T where T is the temporal length for that kernel

    return:
        state_dict: dict, update with inflated weights
    """

    model_scratch_params_dict = dict(model.named_parameters())
    prefix = list(state_dict.keys())[0].split('.')[0]  # grab prefix in the keys for state_dict params

    old_state_dict = copy.deepcopy(state_dict)

    # loop through keys (in either)
    # only check `weights`
    # compare shapes btw 3d model and 2d model
    # if, different, then broadcast
    # then set the broadcasted version into the model value

    for key in sorted(model_scratch_params_dict.keys()):

        scratch_params = model_scratch_params_dict[key]

        # need to add the predix 'model' in convnext
        key_with_prefix = prefix + '.' + key

        # make sure key is in the loaded params first, if not, then print it out
        loaded_params = state_dict.get(key_with_prefix, None)

        if 'time_weight' in key:
            print("found time_weight parameter, train from scratch", key)
            used_params = scratch_params
        elif loaded_params is None:
            # This should never happen for 2D -> 3D ConvNext
            print("Missing key in pretrained model!", key_with_prefix)
            raise Exception
            # used_params = scratch_params

        elif ignore_head and 'head' in key:
            # ignore head weights
            print("found head key / parameter, ignore", key)
            used_params = scratch_params

        elif len(scratch_params.shape) != len(loaded_params.shape):
            # same keys, but inflating weights
            print('key: shape DOES NOT MATCH', key)
            print("scratch:", scratch_params.shape)
            print("pretrain:", loaded_params.shape)
            # need the index [-3], 3rd from last, the temporal dim
            index = -3
            temporal_dim = scratch_params.shape[index]  # temporal len of kernel
            temporal_kernel_factor = 1 if normalize else 1 / temporal_dim
            used_params = repeat(temporal_kernel_factor*loaded_params, '... h w -> ... t h w', t=temporal_dim)
            # loaded_params = temporal_kernel_factor * loaded_params.unsqueeze(index)  # unsqueeze
            # used_params = torch.cat(temporal_dim * [loaded_params], axis=index)  # stack at this dim
        else:
            # print('key: shape MATCH', key)  # loading matched weights
            # used_params = loaded_params
            continue

        state_dict[key_with_prefix] = used_params

    return state_dict

    inputs:
        model: nn.Module, the from 'scratch' model
        state_dict: dict, from the pretrained weights
        ignore_head: bool, whether to inflate weights in the head (or keep scratch weights).
            If number of classes changes (eg, imagenet to hmdb51), then you need to use this.
        normalize: bool, if set to True (default), it inflates with a factor of 1, and if
             set to False it inflates with a factor of 1/T where T is the temporal length for that kernel

    return:
        state_dict: dict, update with inflated weights
    
analyze/convnexts4nd/src/models/baselines/convnext_timm.py