RuntimeError: "param" dim[1] must be greater or equal than dim[2]

[chakrabortym242]

My code is:

import torch
from torch_householder import torch_householder_orgqr

Generate random values

random_tensor = random_tensor = torch.randn(64, 4, 16, 60)

def Household(tensor):
batch_size, num_heads, dim1, dim2 = tensor.shape
reshaped_tensor = tensor.reshape(-1, dim1, dim2) # Shape: [batch * num_heads, dim, dim]
hh= reshaped_tensor.tril(diagonal=-1) + torch.eye(dim1,dim2, device=reshaped_tensor.device)
orthogonal_matrix = torch_householder_orgqr(hh)
# Reshape back to the original shape [batch_size, num_heads, dim, dim]
orthogonal_matrix = orthogonal_matrix.view(batch_size, num_heads, dim1, dim2)
return orthogonal_matrix

It shows error while calculating torch_householder_orgqr