23973f45c189d3cbe34645350cd4f670f7999272,gpytorch/lazy/cat_lazy_tensor.py,CatLazyTensor,_matmul,#CatLazyTensor#Any#,204

Before Change

                rhs = rhs.unsqueeze(0)
            curr_idx = 0
            res_list = []
            index = [slice(None, None, None) for _ in range(self.ndimension())]
            for t, size, rhs in zip(self.lazy_tensors, self.cat_dim_sizes, rhs_):
                index[self.cat_dim] = slice(curr_idx, curr_idx + size, None)
                res_list.append(t._matmul(rhs[index]))
                curr_idx += size

After Change

            res = torch.sum(torch.stack(res_list), dim=0)
        else:
            output_shape = _matmul_broadcast_shape(self.shape, rhs.shape)
            rhs = rhs.expand(*output_shape[:-2], *rhs.shape[-2:])
            curr_idx = 0
            res_list = []
            for t, size in zip(self.lazy_tensors, self.cat_dim_sizes):
                sub_rhs = rhs.narrow(self.cat_dim, curr_idx, size)
                res_list.append(t._matmul(sub_rhs))
                curr_idx += size
            // copy result back to output device
            res_list = [x.to(output_device) for x in res_list]

Instances