f76a4dabb4cd38ee58d01a35c5b511e224d060d2,gpytorch/lazy/sum_batch_lazy_tensor.py,SumBatchLazyTensor,_matmul,#SumBatchLazyTensor#Any#,29

Before Change

        rhs = rhs.unsqueeze(0)
        rhs_size = list(rhs.size())
        rhs_size[0] = self.batch_size()
        rhs = rhs.expand(*rhs_size)

        res = self.base_lazy_tensor._matmul(rhs)
        if self.sum_batch_size is not None:

After Change

        if isvector:
            rhs = rhs.unsqueeze(1)

        if self.num_blocks is None:
            rhs = rhs.unsqueeze(0)
            rhs = rhs.expand(self.base_lazy_tensor.size(0), *tuple(rhs.size())[1:])
        else:
            rhs = rhs.unsqueeze(1)
            rhs = rhs.expand(rhs.size(0), self.num_blocks, *tuple(rhs.size())[2:])
            rhs = rhs.contiguous().view(-1, rhs.size(-2), rhs.size(-1))

        res = self.base_lazy_tensor._matmul(rhs)

        if self.num_blocks is not None:
            res = res.view(-1, self.num_blocks, res.size(1), res.size(2))

Instances