embedding fix

models/transformers.py

@@ -195,7 +195,10 @@ class Transformer_decoder(nn.Module):
         self.decoder_norm = nn.LayerNorm(512)
         self.decoder_layers_parallel = clones(DecoderLayer(512, c(attn), c(attn), c(ff), dropout=0.0), 1)
         self.decoder_norm_parallel = nn.LayerNorm(512)
-        self.cls_embedding = nn.Embedding(52,512)
+        if opts.ref_nshot == 52:
+            self.cls_embedding = nn.Embedding(92,512)
+        else:
+            self.cls_embedding = nn.Embedding(52,512)
         self.cls_token = nn.Parameter(torch.zeros(1, 1, 512))
 
     def forward(self, x, memory, trg_char, src_mask=None, tgt_mask=None):