From 561cec25a4c199e566a784b9bbcabd8e24df5b61 Mon Sep 17 00:00:00 2001
From: Geoffrey Angus <geoffrey@predibase.com>
Date: Wed, 19 Jul 2023 17:30:21 -0700
Subject: [PATCH] pr comments

---
 ludwig/api.py              | 5 ++---
 ludwig/backend/__init__.py | 9 ---------
 ludwig/models/llm.py       | 1 +
 3 files changed, 3 insertions(+), 12 deletions(-)

diff --git a/ludwig/api.py b/ludwig/api.py
index e9b2710265f..16f3e628076 100644
--- a/ludwig/api.py
+++ b/ludwig/api.py
@@ -638,12 +638,13 @@ def on_epoch_end(self, trainer, progress_tracker, save_path):
                         test_set=test_set,
                         save_path=model_dir,
                     )
+                    (self.model, train_trainset_stats, train_valiset_stats, train_testset_stats) = train_stats
 
                     # Calibrates output feature probabilities on validation set if calibration is enabled.
                     # Must be done after training, and before final model parameters are saved.
                     if self.backend.is_coordinator():
                         calibrator = Calibrator(
-                            trainer.model,
+                            self.model,
                             self.backend,
                             batch_size=trainer.eval_batch_size,
                         )
@@ -667,7 +668,6 @@ def on_epoch_end(self, trainer, progress_tracker, save_path):
                         if not skip_save_model:
                             # ensure that any changes to the model object held by the
                             # trainer class are reflected in the model in this class.
-                            self.model = trainer.model
                             self.model.save(model_dir)
 
                     # Evaluation Frequency
@@ -687,7 +687,6 @@ def on_epoch_end(self, trainer, progress_tracker, save_path):
                     # List[TrainerMetric], with one entry per training checkpoint, according to steps_per_checkpoint.
                     # We reduce the dictionary of TrainerMetrics to a simple list of floats for interfacing with Ray
                     # Tune.
-                    (self.model, train_trainset_stats, train_valiset_stats, train_testset_stats) = train_stats
                     train_stats = TrainingStats(
                         metric_utils.reduce_trainer_metrics_dict(train_trainset_stats),
                         metric_utils.reduce_trainer_metrics_dict(train_valiset_stats),
diff --git a/ludwig/backend/__init__.py b/ludwig/backend/__init__.py
index a088a43d499..d5b84f934d1 100644
--- a/ludwig/backend/__init__.py
+++ b/ludwig/backend/__init__.py
@@ -59,15 +59,6 @@ def _has_ray():
         return False
 
 
-def is_ray_backend(backend) -> bool:
-    if isinstance(backend, str):
-        return backend == "ray"
-    elif isinstance(backend, dict):
-        return backend.get("type", "local") == "ray"
-    else:
-        return False
-
-
 def get_local_backend(**kwargs):
     return LocalBackend(**kwargs)
 
diff --git a/ludwig/models/llm.py b/ludwig/models/llm.py
index 8c8cb705d1c..06c25c19693 100644
--- a/ludwig/models/llm.py
+++ b/ludwig/models/llm.py
@@ -194,6 +194,7 @@ def initialize_adapter(self):
             logger.info("==================================================")
 
     def prepare_for_training(self):
+        # TODO: this implementation will not work if resuming from a previous checkpoint. Need to fix this.
         self.initialize_adapter()
 
     def to_device(self, device):