Update rita_modeling.py

rita_modeling.py

@@ -217,15 +217,15 @@ class RITAModel(PreTrainedModel):
         config
     ):
         super().__init__(config)
-        self.embedding = nn.Embedding(config.in_vocab_size, config.d_model)
+        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
         self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_layers)])
         self.final_norm = nn.LayerNorm(config.d_model)
-        self.projector = nn.Linear(config.d_model, config.out_vocab_size, bias = False)
+        self.projector = nn.Linear(config.d_model, config.vocab_size, bias = False)
 
     def forward(self, ids, attn_mask=None, padding_mask=None, return_hidden=False) -> torch.FloatTensor:
         x = self.embedding(ids)  # N x L x D
         if attn_mask == None:
-            attn_mask = (torch.triu(torch.ones(ids.size(1), ids.size(1))) == 0).transpose(0, 1).contiguous()
+            attn_mask = (torch.triu(torch.ones(ids.size(1), ids.size(1))) == 0).transpose(0, 1).contiguous().to(ids.device)
         for layer in self.layers:
             x = layer(x, attn_mask=attn_mask, padding_mask=padding_mask)
         x = self.final_norm(x)  # N x L x D