Update DistilBERT.py

DistilBERT.py

@@ -89,6 +89,10 @@ class DistilBERTClass(torch.nn.Module):
         self.classifier = torch.nn.Linear(768, 1)
 
     def forward(self, input_ids, attention_mask, token_type_ids):
+        # Convert input_ids to PyTorch tensor if it's a list
+        if isinstance(input_ids, list):
+            input_ids = torch.tensor(input_ids)
+
         output_1 = self.l1(input_ids=input_ids, attention_mask=attention_mask)
         hidden_state = output_1[0]
         pooler = hidden_state[:, 0]
@@ -98,8 +102,6 @@ class DistilBERTClass(torch.nn.Module):
         output = self.classifier(pooler)
         return output
 
-model_DB = DistilBERTClass()
-model_DB.to(device)
 
 # Validation function
 def validation(testing_loader):