Upload main.py with huggingface_hub

main.py

@@ -253,9 +253,11 @@ criterion = nn.SmoothL1Loss(reduction="none")
 
 num_epochs = train_config["num_epochs"]
 
+# Only pass trainable parameters to DeepSpeed (frozen params cause grad tracking errors)
+trainable_params = [p for p in model.parameters() if p.requires_grad]
 model_engine, optimizer, _, _ = deepspeed.initialize(args=args,
                                                      model=model,
-                                                     model_parameters=model.parameters(),
+                                                     model_parameters=trainable_params,
                                                      )
 
 global_rank = deepspeed.comm.get_rank()