[Need Help] TypeError: reshape total size must be unchanged

[Yuji-github]

I tried to update params that involve a pre-trained model, however, I get mismatch dims.
TypeError: reshape total size must be unchanged, got new_sizes (1, 1, 16, 64) for shape (1, 1, 256, 1024)

I think I could not pass the network's variables in the 'params' key properly.
Here is what I have done for my model:

    class MyNLP(nn.Module):
        roberta: nn.Module

        @nn.compact
        def __call__(self, input_ids, attention_mask):  
            out = self.roberta(input_ids, attention_mask, output_hidden_states=True)  
            cls_embedding = out.hidden_states[0][:,0,:]       
            out = nn.Dense(features=3, name="Prediction")(cls_embedding)
    
            return out

    init_rng = jax.random.PRNGKey(0)
    model = MyNLP(pretrained_roberta)
    variables = model.init(init_rng, dummy_input_ids, dummy_attention_mask)  # <- store only Prediction variables
  
    variables = unfreeze(variables)  
    variables['params']['RoberTa'] = pretrained_roberta.params  # <- add roberta variables to the Prediction variables
    variables = freeze(variables)
  
    opt = optax.adam(learning_rate=2e-5, b1=0.9, b2=0.999)
    
    state = TrainState.create(apply_fn=model.apply,  params=variables['params'],  tx=opt)

    params = state.params

    def data_loader(dataset, batch_size): 
        for i in range(len(dataset) // batch_size):
            batch = dataset[i * batch_size : (i + 1) * batch_size]
            batch = {k: jnp.array(v) for k, v in batch.items()}          
            yield batch

    for batch in data_loader(data, total_batch_size):  
        labels = batch.pop("labels"):
        logits = state.apply_fn({'params': state.params}, **batch)  # <- ERROR HERE 

How can I fix the errors in this case?

[marcvanzee]

Some comments:

• I edited your comment to make the code use syntax highlighting for Python.
• If you don't do any training you do not need to use an optimizer or a train state. You can simply call model.apply directly.

Question: What is pretrained_roberta? I suppose it does some reshape and your input data does not have the right dimensions for this reshape.