Update archctr.py

archctr.py

@@ -1,3 +1,6 @@
+import torch
+import torch.nn as nn
+
 class GPTDecoder(nn.Module):
     def __init__(self, nclass=16135, d_model=1024, max_pos=512, nhead=16, dim_feedforward=2048, vocab_size=16135, nhead_emb=8):
         super().__init__()
@@ -16,7 +19,7 @@ class GPTDecoder(nn.Module):
         self.Linear = nn.Linear(d_model, nclass)
     def forward(self, x):
         x = self.embedding(x)
-        x_n = torch.arange(x.size(1)).unsqueeze(0).to('cuda')
+        x_n = torch.arange(x.size(1)).unsqueeze(0)
         x_n = self.pos_embedding(x_n)
         x = x + x_n
         x_origin = x
@@ -28,7 +31,7 @@ class GPTDecoder(nn.Module):
         x = x + x_origin
         x_origin = x
         x = self.LayerNorm_heads(x)
-        mask = nn.Transformer.generate_square_subsequent_mask(x.size(1)).to('cuda')
+        mask = nn.Transformer.generate_square_subsequent_mask(x.size(1))
         x, _ = self.GPT_MHA(x,x,x, attn_mask=mask)
         x = x+x_origin
         x_origin = x