This commit is contained in:
Daniel Povey 2022-05-18 10:26:29 +08:00
parent 668d01cc7a
commit ebc2ffeff7
2 changed files with 2 additions and 2 deletions

View File

@ -632,7 +632,7 @@ class Cain(Optimizer):
rev_dims_order.append(ndim-1)
for i in range(dim+1, ndim):
dims_order.append(i)
rev_dims_order.append(i)
rev_dims_order.append(i-1)
dims_order.append(dim)
# e.g. ndim=4, dim=1, dims_order=(0,2,3,1), rev_dims_order=(0,3,1,2)
new_grad = grad.permute(*dims_order)

View File

@ -628,7 +628,7 @@ class Eve(Optimizer):
rev_dims_order.append(ndim-1)
for i in range(dim+1, ndim):
dims_order.append(i)
rev_dims_order.append(i)
rev_dims_order.append(i-1)
dims_order.append(dim)
# e.g. ndim=4, dim=1, dims_order=(0,2,3,1), rev_dims_order=(0,3,1,2)
new_grad = grad.permute(*dims_order)