Update inference.py

inference.py

@@ -420,6 +420,9 @@ class GenerativeInferenceModel:
                     for source_key, value in state_dict.items():
                         if source_key.startswith('module.model.'):
                             target_key = source_key[len('module.model.'):]
+                            # Some ckpts have 'module.model.model.<...>'; remove the extra 'model.' too
+                            if target_key.startswith('model.'):
+                                target_key = target_key[len('model.'):]
                             resnet_state_dict[target_key] = value
                             
                     print(f"Extracted {len(resnet_state_dict)} parameters from module.model")