in TransformerXLLayer code snippet will face with issues

[alokamgnaneswarasai]

def forward(self, *, x: torch.Tensor, mem: Optional[torch.Tensor], mask: torch.Tensor): """ *xis a tensor of the token level feature vectors of shape[seq_len, batch_size, d_model]*memis a tensor of the past token level feature vectors of shape[mem_len, batch_size, d_model]*maskis a matrix of shape[seq_len, mem_len + seq_len, batch_size]or[seq_len, mem_len + seq_len, 1]. mask[i, j]is true if token atican see token atj. """ # Normalize the vectors before doing self attention z = self.norm_self_attn(x) # If there is memory if mem is not None: # Normalize it mem = self.norm_self_attn(mem) # Concatenate with zm_z = torch.cat((mem, z), dim=0) # Ignore if there is no memory

In this code snippet what happens if the current input batch size does not match with the previous input batch size ??. Lets say total number of training samples is 68 and batch size is 89 then last batch has only 4 samples , which does not matches with batch size of 8 , then how will you concatenate mem,z which are of different shapes