wip: cleanup

src/ntl/trainer.py

@@ -1297,9 +1297,7 @@ def compute_loss(self, model, inputs, return_outputs=False):
             labels = inputs["labels"]
         else:
             labels = None
-        print('Before IN', type(model))
         outputs = model(**inputs)
-        print('Done with forward pass')
         # Save past state if it exists
         # TODO: this needs to be fixed and made cleaner later.
         if self.args.past_index >= 0:

src/ntl/transformer_backbone/t5/t5_vanilla_for_number_token_loss.py

@@ -59,7 +59,6 @@ def forward(
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
         )
-        print("TYPE outputs", type(outputs))
 
         # If labels are provided, calculate and combine the NumberTokenLoss
         if labels is not None and self.number_token_loss is not None:
@@ -70,6 +69,5 @@ def forward(
                 token_loss=outputs.loss,
                 loss=(outputs.loss + self.number_token_loss.weight * number_token_loss),
             )
-            print("type", type(outputs), outputs.loss, outputs.number_loss)
 
         return outputs