Upload train_main_model.py with huggingface_hub

train_main_model.py

@@ -169,7 +169,7 @@ def train_with_best_params(
         # Training
         color_model = feature_models[config.color_column]
         hierarchy_model = feature_models[config.hierarchy_column]
-        train_loss, align_metrics = train_one_epoch_enhanced(
+        train_loss, align_metrics = train_one_epoch(
             model, train_loader, optimizer, feature_models, color_model, hierarchy_model,
             config.device, processor, temperature, alignment_weight,
             reference_model=reference_clip, reference_weight=0.1