Ensure matrix dims are correct after matmul
@@ -4,16 +4,12 @@ from __future__ import print_function
from __future__ import unicode_literals
import torch
-from IPython.core.debugger import Pdb
import operator
from torch.autograd import Variable
from .lazy_variable import LazyVariable
from functools import reduce
-ipdb = Pdb()
-
-
def _prod(iterable):
return reduce(operator.mul, iterable, 1)
@@ -26,12 +22,15 @@ def _matmul(lazy_vars, rhs):
if is_batch:
n_batch = res.size(0)
res = res.transpose(-2, -1).contiguous().view(n_batch, n_cols, lazy_var.size(-1), -1).transpose(0, 1).contiguous()
- # ipdb.set_trace()
factor = lazy_var._matmul(res).permute(1, 3, 2, 0)
res = factor.contiguous().view(n_batch, -1, n_cols)
else:
res = res.t().contiguous().view(n_cols, lazy_var.size(-1), -1)
- factor = lazy_var._matmul(res).permute(2, 1, 0)
+ factor = lazy_var._matmul(res)
+ if factor.ndimension() == 3:
+ factor = factor.permute(2, 1, 0)
+ else:
+ factor = factor.transpose(-2, -1)
res = factor.contiguous().view(-1, n_cols)
return res
