diff --git a/stanza/models/lemma_classifier/evaluate_models.py b/stanza/models/lemma_classifier/evaluate_models.py
index e5a8692fa3..fc041345ce 100644
--- a/stanza/models/lemma_classifier/evaluate_models.py
+++ b/stanza/models/lemma_classifier/evaluate_models.py
@@ -153,12 +153,6 @@ def evaluate_model(model: nn.Module, eval_path: str, verbose: bool = True, is_tr
 
     # load in eval data
     text_batches, index_batches, label_batches, _, label_decoder = utils.load_dataset(eval_path, label_decoder=model.label_decoder)
-
-    # TODO fix this in the future
-    text_batches, index_batches, label_batches = text_batches[: -1], index_batches[: -1], label_batches[: -1]
-    
-    index_batches = torch.stack(index_batches).to(device)
-    label_batches = torch.stack(label_batches).to(device)
     
     logging.info(f"Evaluating on evaluation file {eval_path}")
 
diff --git a/stanza/models/lemma_classifier/train_model.py b/stanza/models/lemma_classifier/train_model.py
index 5fad39ef42..37185c1473 100644
--- a/stanza/models/lemma_classifier/train_model.py
+++ b/stanza/models/lemma_classifier/train_model.py
@@ -124,10 +124,6 @@ def train(self, num_epochs: int, save_name: str, args: Mapping, eval_file: str,
             logging.info(f"Loaded dataset successfully from {train_path}")
             logging.info(f"Using label decoder: {label_decoder}  Output dimension: {self.output_dim}")
 
-            text_batches, idx_batches, label_batches = text_batches[:-1], idx_batches[:-1], label_batches[:-1]  # TODO come up with a fix for this
-            
-            idx_batches, label_batches = torch.stack(idx_batches).to(device), torch.stack(label_batches).to(device)  
-
         self.model = LemmaClassifierLSTM(self.vocab_size, self.embedding_dim, self.hidden_dim, self.output_dim, self.vocab_map, self.embeddings, label_decoder,
                                          charlm=self.use_charlm, charlm_forward_file=self.forward_charlm_file, charlm_backward_file=self.backward_charlm_file)
         self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)
diff --git a/stanza/models/lemma_classifier/transformer_baseline/baseline_trainer.py b/stanza/models/lemma_classifier/transformer_baseline/baseline_trainer.py
index a7d51b75a7..0f552709a3 100644
--- a/stanza/models/lemma_classifier/transformer_baseline/baseline_trainer.py
+++ b/stanza/models/lemma_classifier/transformer_baseline/baseline_trainer.py
@@ -108,18 +108,12 @@ def train(self, num_epochs: int, save_name: str, args: Mapping, eval_file: str,
             text_batches, position_batches, label_batches, counts, label_decoder = utils.load_dataset(kwargs.get("train_path"), get_counts=self.weighted_loss)
             self.output_dim = len(label_decoder)
             logging.info(f"Using label decoder : {label_decoder}")
-
-            # # TODO: fix this to make it not disregard last batch, and instead pad it or some other idea
-            # text_batches, position_batches, label_batches = text_batches[:-1], position_batches[:-1], label_batches[:-1]
-
-            # # Move data to device
-            # label_batches = torch.stack(label_batches).to(device)
-            # position_batches = torch.stack(position_batches).to(device)
         
         assert len(text_batches) == len(position_batches) == len(label_batches), f"Input batch sizes did not match ({len(text_batches)}, {len(position_batches)}, {len(label_batches)})."
 
         self.model = LemmaClassifierWithTransformer(output_dim=self.output_dim, transformer_name=self.transformer_name, label_decoder=label_decoder)
-        self.optimizer = self.set_layer_learning_rates(transformer_lr=self.lr/2, mlp_lr=self.lr)  # Adam optimizer
+        # self.optimizer = self.set_layer_learning_rates(transformer_lr=self.lr/2, mlp_lr=self.lr)  # Adam optimizer
+        self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)
 
         self.model.to(device)
         self.model.transformer.to(device)