Although I will mention that a return type here would be similar to how PyTorch handles docs for in place changes.
@patch
def dihedral(x:TensorImage,
k:int, # Dihedral transformation to apply
) -> TensorImage:
if k in [1,3,4,7]: x = x.flip(-1)
if k in [2,4,5,7]: x = x.flip(-2)
if k in [3,5,6,7]: x = x.transpose(-1,-2)
return x
Also removed the addition return docment, per our discussions.
Although I will mention that a return type here would be similar to how PyTorch handles docs for in place changes.
Also removed the addition return docment, per our discussions.