diff --git a/src/bert_training.py b/src/bert_training.py
index e227de3..ec78b55 100644
--- a/src/bert_training.py
+++ b/src/bert_training.py
@@ -147,14 +147,13 @@ class Encoder(nn.Module):
         self.num_shifts = nn.Parameter(torch.randn(1, len(self.num_columns)))
         self.proj = nn.Linear(self.total_h_dim, self.out_dim, bias=False)
         
-    def forward(self, cat_features_batch, num_features_batch, targets_batch):
+    def forward(self, cat_features_batch, num_features_batch):
         cat_embed_tensor = self.cat_embeds(cat_features_batch.type(torch.int32))
         cat_embed_tensor = cat_embed_tensor.reshape(cat_features_batch.shape[0], cat_features_batch.shape[1], -1)
         num_embed_tensor = self.num_scales * num_features_batch + self.num_shifts        
         embed_tensor = torch.concat([cat_embed_tensor, num_embed_tensor], dim=-1)
         inputs = self.proj(embed_tensor)
-        targets = targets_batch
-        return inputs, targets
+        return inputs
 
 # RoFormer: Enhanced Transformer with Rotary Position Embedding https://arxiv.org/abs/2104.09864
 class RoPE(nn.Module):
@@ -247,9 +246,9 @@ class Model(nn.Module):
         self.encoder = encoder
         self.classifier = classifier
         
-    def forward(self, cat_inputs, num_inputs, targets):
-        inputs, targets = self.encoder(cat_inputs, num_inputs, targets)
-        return self.classifier(inputs), targets
+    def forward(self, cat_inputs, num_inputs):
+        inputs = self.encoder(cat_inputs, num_inputs)
+        return self.classifier(inputs)
 
 def test(start_time, epoch, batches_per_epoch, batch_size, model, optimizer, credit_dataset, test_auroc, writer):
         model.eval()
@@ -260,8 +259,8 @@ def test(start_time, epoch, batches_per_epoch, batch_size, model, optimizer, cre
                 test_cat_inputs = test_cat_inputs.to("cuda", non_blocking=True)
                 test_num_inputs = test_num_inputs.to("cuda", non_blocking=True)
                 test_targets = test_targets.to("cuda", non_blocking=True)
-                outputs, targets = model(test_cat_inputs, test_num_inputs, test_targets)
-                test_auroc.update(outputs, targets.long())
+                outputs = model(test_cat_inputs, test_num_inputs)
+                test_auroc.update(outputs, test_targets.long())
                 print(f"\r {test_batch_id}/{len(credit_dataset.test_uniq_client_ids)//batch_size} {test_auroc.compute().item():.5f}", end = " "*20)
         writer.add_scalar('test_roc_auc', test_auroc.compute().item(), epoch * batches_per_epoch)
         print(f"\r {datetime.now() - start_time} {epoch}/{epochs} Test rocauc: {test_auroc.compute().item():.5f}", end = " "*20) 
@@ -361,12 +360,11 @@ try:
             model.train()
             optimizer.train()
             optimizer.zero_grad()
-            outputs, targets = model(
+            outputs = model(
                 cat_inputs[0].to("cuda"), 
-                num_inputs[0].to("cuda"), 
-                targets[0].to("cuda")
+                num_inputs[0].to("cuda")
             )
-            loss = criterion(outputs, targets)
+            loss = criterion(outputs, targets[0].to("cuda"))
             loss.backward()
             optimizer.step()