حذف حلقه for

پروژه ی قبلی آقای محمد نصیری سلام من چجوری میتونم حلقه ی for رو توی این کد با استفاده از خصوصیات pytorch حذفش کنم: import torch length = 3 def WToSN(In):    LenSN = 2 \*\* length    SN = torch.zeros(LenSN, device=torch.device('cuda'))    SN\[0:In\] = 1    return SNweight = torch.tensor(\[\[1, 2, 3\], \[4, 5, 6\]\], device=torch.device('cuda'))print(weight)WeightSN = torch.stack(\[WToSN(x) for x in weight.flatten()\]).reshape(weight.size(0),weight.size(1), -1)print(WeightSN)