CausalVAE u纬度和A不同时应该如何修改，直接将u通过mlp映射到相同的纬度可以吗

[AndrewTal]

def mask_u(self,x):
    self.B = self.A
    #if self.i:
    #    x = x.view(-1, x.size()[1], 1)
    #    x = torch.matmul((self.B+0.5).t().int().float(), x)
    #    return x
    x = x.view(-1, x.size()[1], 1)
    x = torch.matmul(self.B.t(), x)
    return x

您好，如果 u 和 causal 矩阵 A 的维度不同，例如只 flow 数据集中 4 个参数只给出其中 1 个作为 u，代码应该如何修改？直接将u通过mlp映射到A相同的纬度可以吗