Error fixes

[QuantumSorcerer02]

--- FIX FOR gemma/model.py ---

def load_weights(model, weights_path):
state_dict = torch.load(weights_path)
new_state_dict = {}

for key, value in state_dict.items():
    # FIX: Check if the weight uses the new Gemma 4 multimodal prefix
    if key.startswith("model.language_model."):
        # Map back to the flat structure expected by the text-only class
        new_key = key.replace("model.language_model.", "model.")
        new_state_dict[new_key] = value
    else:
        new_state_dict[key] = value
        
model.load_state_dict(new_state_dict)

[sonali-kumari1]

Hi @QuantumSorcerer02 -

Thanks for sharing this. Could you share exact Gemma PyTorch version/commit you’re using, along with a minimal reproducible example showing the full error traceback?