Update model.py

model.py

@@ -2,15 +2,11 @@ import torch
 import torchtext
 from torchtext.models import RobertaClassificationHead, XLMR_BASE_ENCODER
 from torch import nn
-def xlmr_base_encoder_model(num_classes:int=2, # default output classes = 2 (Bad, Good)):
-  # 1, 2, 3 Create EffNetB2 pretrained weights, transforms and model
-  transforms = torchtext.models.XLMR_BASE_ENCODER.transform()
-  classifier_head = RobertaClassificationHead(num_classes = 2, input_dim = 768)
-  model = XLMR_BASE_ENCODER.get_model(head = classifier_head)
-
-  # 4. Freeze all layers in the base model
-  for param in model.parameters():
-    param.requires_grad = False
-
-
-  return model, transforms
+def xlmr_base_encoder_model(num_classes:int=2, # default output classes = 2 (Bad, Good)
+                           ):
+                               transforms = XLMR_BASE_ENCODER.transform()
+                               classifier_head = RobertaClassificationHead(num_classes = 2, input_dim = 768)
+                               model = XLMR_BASE_ENCODER.get_model(head = classifier_head)
+                               for param in model.parameters():
+                                   param.requires_grad = False
+                                return model, transforms