ccc913a65e08bc5523eb2490d3177c472b55d094,gpytorch/lazy/sum_batch_lazy_variable.py,SumBatchLazyVariable,__getitem__,#SumBatchLazyVariable#Any#,127

Before Change

    def __getitem__(self, index):
        if self.sum_batch_size is None:
            return super(SumBatchLazyVariable, self).__getitem__(index)

        // Cases for when there"s an inner batch
        else:
            batch_index = index if isinstance(index, int) else index[0]

After Change

            if isinstance(index, int):
                index = tuple(slice(None, None, None), int)
            else:
                index = tuple([slice(None, None, None)] + list(index))

            // Do a __getitem__ on the base lazy variable
            res = self.base_lazy_variable.__getitem__(index)
            if isinstance(res, LazyVariable):
                return SumBatchLazyVariable(res, sum_batch_size=None)
            else:
                return res.sum(0)

        // Cases for when there"s an inner batch
        else:
            batch_index = index if isinstance(index, int) else index[0]

            // Keeping all batch dimensions - recursion base case

Instances