Fix get indices for kronecker
@@ -154,7 +154,7 @@ class KroneckerProductLazyVariable(LazyVariable):
return self.__class__(*(lazy_var._transpose_nonbatch() for lazy_var in self.lazy_vars), **self._kwargs)
def _batch_get_indices(self, batch_indices, left_indices, right_indices):
- res = Variable(self._tensor_cls(left_indices.size()).fill_(1))
+ res = Variable(self.tensor_cls(left_indices.size()).fill_(1))
size = self.size(-1)
for i, lazy_var in enumerate(list(self.lazy_vars)[::-1]):
size = size / lazy_var.size(-1)
@@ -167,7 +167,7 @@ class KroneckerProductLazyVariable(LazyVariable):
return res
def _get_indices(self, left_indices, right_indices):
- res = Variable(self._tensor_cls(left_indices.size()).fill_(1))
+ res = Variable(self.tensor_cls(left_indices.size()).fill_(1))
size = self.size(-1)
for i, lazy_var in enumerate(list(self.lazy_vars)[::-1]):
size = size / lazy_var.size(-1)
