def backward(self, grad_output):
means, = self.saved_tensors
T = self.R.shape[0]
dim = means.dim()
# Add batch axis if necessary
if dim == 2:
T_, D = means.shape
B = 1
grad_output = grad_output.view(B, T, -1)
else:
B, T_, D = means.shape
grad = torch.matmul(self.R.transpose(0, 1), grad_output)
reshaped = not (T == T_)
if not reshaped:
grad = grad.view(B, self.num_windows, T, -1).transpose(
1, 2).contiguous().view(B, T, D)
if dim == 2:
return grad.view(-1, D)
return grad
评论列表
文章目录