blas.py 文件源码

def backward(self, grad_output):
        batch1, batch2 = self.saved_tensors
        grad_add_batch = grad_batch1 = grad_batch2 = None

        if self.needs_input_grad[0]:
            grad_add_batch = grad_output
            if self.alpha != 1:
                grad_add_batch = grad_add_batch.mul(self.alpha)

        if self.needs_input_grad[1]:
            grad_batch1 = torch.bmm(grad_output, batch2.transpose(1, 2))
            if self.beta != 1:
                grad_batch1 *= self.beta

        if self.needs_input_grad[2]:
            grad_batch2 = torch.bmm(batch1.transpose(1, 2), grad_output)
            if self.beta != 1:
                grad_batch2 *= self.beta

        return grad_add_batch, grad_batch1, grad_batch2