Remove torchinfo

app.py

@@ -17,7 +17,7 @@ with open("example_texts.txt", "r") as file:
 
 ### Model and transforms preparation ###
 # Create model and tokenizer
-model, tokenizer = create_roberta_model(output_shape=len(class_names), print_summary=False)
+model, tokenizer = create_roberta_model(output_shape=len(class_names))
 
 # Load saved weights
 model.load_state_dict(

model.py

@@ -1,5 +1,4 @@
 import torch
-from torchinfo import summary
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -9,7 +8,7 @@ def model_input_wrapper(batch_size, sequence_length, tokenizer):
   dummy_attention_mask = torch.ones(batch_size, sequence_length, dtype=torch.long)
   return {'input_ids': dummy_input_ids, 'attention_mask': dummy_attention_mask}
 
-def create_roberta_model(output_shape:int=10, device=device, print_summary=True):
+def create_roberta_model(output_shape:int=10, device=device):
   """Creates a HuggingFace roberta-base model.
 
   Args:
@@ -31,8 +30,4 @@ def create_roberta_model(output_shape:int=10, device=device, print_summary=True)
 
   model.classifier.out_proj = torch.nn.Linear(in_features=768, out_features=output_shape)
 
-  if print_summary:
-    sample_inputs = model_input_wrapper(1, 128, tokenizer)
-    print(summary(model, input_data=sample_inputs, verbose=0, col_names=["input_size", "output_size", "num_params", "trainable"], col_width=20, row_settings=["var_names"]))
-
   return model.to(device), tokenizer