def stacklists(arg):
 if isinstance(arg, (tuple, list)):
  return stack(list(map(stacklists, arg)))
 else:
  return arg

def stacklists ( arg ) : if isinstance ( arg , ( tuple , list ) ) : return stack ( list ( map ( stacklists , arg ) ) ) else : return arg
recursively stack lists of tensors to maintain similar structure .