removing next call to fix StopIteration error

This commit is contained in:
TevenLeScao
2020-11-02 14:04:28 +01:00
parent 84caa23301
commit 7d108a92f2
+5 -5
View File
@@ -330,14 +330,14 @@ class RelPartialLearnableMultiHeadAttn(nn.Module):
if attn_mask is not None and torch.sum(attn_mask).item():
attn_mask = attn_mask == 1 # Switch to bool
if attn_mask.dim() == 2:
if next(self.parameters()).dtype == torch.float16:
if list(self.parameters())[0].dtype == torch.float16:
attn_score = (
attn_score.float().masked_fill(attn_mask[None, :, :, None], -65000).type_as(attn_score)
)
else:
attn_score = attn_score.float().masked_fill(attn_mask[None, :, :, None], -1e30).type_as(attn_score)
elif attn_mask.dim() == 3:
if next(self.parameters()).dtype == torch.float16:
if list(self.parameters())[0].dtype == torch.float16:
attn_score = attn_score.float().masked_fill(attn_mask[:, :, :, None], -65000).type_as(attn_score)
else:
attn_score = attn_score.float().masked_fill(attn_mask[:, :, :, None], -1e30).type_as(attn_score)
@@ -435,7 +435,7 @@ class AdaptiveEmbedding(nn.Module):
if self.d_proj != self.d_embed:
embed = F.linear(embed, self.emb_projs[0])
else:
param = next(self.parameters())
param = list(self.parameters())[0]
inp_flat = inp.view(-1)
emb_flat = torch.zeros([inp_flat.size(0), self.d_proj], dtype=param.dtype, device=param.device)
for i in range(len(self.cutoffs)):
@@ -806,7 +806,7 @@ class TransfoXLModel(TransfoXLPreTrainedModel):
def init_mems(self, bsz):
if self.mem_len > 0:
mems = []
param = next(self.parameters())
param = list(self.parameters())[0]
for i in range(self.n_layer):
empty = torch.zeros(self.mem_len, bsz, self.config.d_model, dtype=param.dtype, device=param.device)
mems.append(empty)
@@ -886,7 +886,7 @@ class TransfoXLModel(TransfoXLPreTrainedModel):
elif head_mask.dim() == 2:
head_mask = head_mask.unsqueeze(1).unsqueeze(1).unsqueeze(1)
head_mask = head_mask.to(
dtype=next(self.parameters()).dtype
dtype=list(self.parameters())[0].dtype
) # switch to fload if need + fp16 compatibility
else:
head_mask = [None] * self.n_layer