Fix indexing bug in LazyEvaluatedKernelTensor
Seems to have been a copy&paste error
@@ -113,12 +113,12 @@ class LazyEvaluatedKernelTensor(LazyTensor):
# We're going to handle multi-batch indexing with a try-catch loop
# This way - in the default case, we can avoid doing expansions of x1 which can be timely
except IndexError:
- if any([not isinstance(bi, slice) for bi in batch_indices]):
+ if any(not isinstance(bi, slice) for bi in batch_indices):
raise RuntimeError(
"Attempting to tensor index a non-batch matrix's batch dimensions. "
f"Got batch index {batch_indices} but my shape was {self.shape}"
)
- x2 = x2.expand(*([1] * (len(batch_indices) - self.x1.dim() + 2)), *self.x2.shape)
+ x2 = x2.expand(*([1] * (len(batch_indices) - self.x2.dim() + 2)), *self.x2.shape)
x2 = x2[(*batch_indices, col_index, dim_index)]
if len(batch_indices) == 0 or all(ind == slice(None, None, None) for ind in batch_indices):
