Update model.py

model.py

@@ -5,7 +5,7 @@ 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 = torchtext.RobertaClassificationHead(num_classes = 2, input_dim = 768)
+  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