fix training loop for push to hub

training_loop.py

@@ -74,12 +74,14 @@ def main(
 
     if push_to_hub:
         config = AutoConfig.from_pretrained(f"nvidia/mit-b{model_flavor}")
-        config.push_to_hub("segformer-sidewalk", repo_url="https://huggingface.co/ChainYo/segformer-sidewalk")
+        config.num_labels = num_labels
+        config.id2label = id2label
+        config.label2id = {v: k for k, v in id2label_file.items()}
+        config.push_to_hub(".", repo_url="https://huggingface.co/ChainYo/segformer-sidewalk")
+
         checkpoint_path = checkpoint_callback.best_model_filepath
-        model = SegformerForSemanticSegmentation.from_pretrained(
-            checkpoint_path, num_labels=num_labels, id2label=id2label, label2id=id2label, config=config,
-        )
-        model.push_to_hub("segformer-sidewalk", repo_url="https://huggingface.co/ChainYo/segformer-sidewalk")
+        model = SegformerForSemanticSegmentation.from_pretrained(checkpoint_path, config=config,)
+        model.push_to_hub(".", repo_url="https://huggingface.co/ChainYo/segformer-sidewalk")
 
 
 if __name__ == "__main__":