Update StreamlitModel.py

StreamlitModel.py

@@ -11,7 +11,6 @@ Created on Mon Feb 19 13:56:57 2024
 #For more info please refer to the Github script
 
 import torch
-from torch import nn
 from transformers import GPT2Tokenizer, GPT2Config, GPT2Model, GPT2PreTrainedModel
 from torch.nn import functional as F
 from better_profanity import profanity
@@ -34,7 +33,7 @@ class GPT2_Model(GPT2PreTrainedModel):
         # this is necessary since we add a new unique token for pad_token
         self.transformer.resize_token_embeddings(len(tokenizer))
 
-        self.lm_head = nn.Linear(config.n_embd, len(tokenizer), bias=False)
+        self.lm_head = torch.nn.Linear(config.n_embd, len(tokenizer), bias=False)
 
     def forward(self, input_ids, attention_mask=None, token_type_ids=None):