Check failed: rank <= DML_TENSOR_DIMENSION_COUNT_MAX

Hi, i just updated to version 0.2.0.dev230426 (with previous versions 0.1.13.1.dev230413 there was not this problem) and now it shows this error: Error: [F D:\a_work\1\s\pytorch-directml\caffe2\dml\dml_tensor_desc.cc:81] Check failed: rank <= DML_TENSOR_DIMENSION_COUNT_MAX

Windows 10 Python 3.10.10 torch-directml 0.2.0.dev230426 AMD Rx 6600

Code:


class IFNet(nn.Module):
def __init__(self, backend):
    super(IFNet, self).__init__()
    self.block0 = IFBlock(7+4, c=90)
    self.block1 = IFBlock(7+4, c=90)
    self.block2 = IFBlock(7+4, c=90)
    self.block_tea = IFBlock(10+4, c=90)

    self.backend = backend

def forward(self, x, scale_list=[4, 2, 1], training=False):
    if training == False:
        channel = x.shape[1] // 2
        img0 = x[:, :channel]
        img1 = x[:, channel:]
    flow_list = []
    merged = []
    mask_list = []
    warped_img0 = img0
    warped_img1 = img1
    flow = (x[:, :4]).detach() * 0
    mask = (x[:, :1]).detach() * 0
    loss_cons = 0
    block = [self.block0, self.block1, self.block2]
    for i in range(3):
        f0, m0 = block[i](torch.cat((warped_img0[:, :3], warped_img1[:, :3], mask), 1), flow, scale=scale_list[i])
        f1, m1 = block[i](torch.cat((warped_img1[:, :3], warped_img0[:, :3], -mask), 1), torch.cat((flow[:, 2:4], flow[:, :2]), 1), scale=scale_list[i])
        flow = flow + (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2
        mask = mask + (m0 + (-m1)) / 2
        mask_list.append(mask)
        flow_list.append(flow)
        warped_img0 = warp(img0, flow[:, :2], self.backend)
        warped_img1 = warp(img1, flow[:, 2:4], self.backend)
        merged.append((warped_img0, warped_img1))
    '''
    c0 = self.contextnet(img0, flow[:, :2])
    c1 = self.contextnet(img1, flow[:, 2:4])
    tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1)
    res = tmp[:, 1:4] * 2 - 1
    '''
    for i in range(3):
        mask_list[i] = torch.sigmoid(mask_list[i])
        merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i])
        # merged[i] = torch.clamp(merged[i] + res, 0, 1)        
    return flow_list, mask_list[2], merged

class RIFEv3: def init(self, backend, local_rank=-1): self.flownet = IFNet(backend) self.device(backend) self.optimG = AdamW(self.flownet.parameters(), lr=1e-6, weight_decay=1e-4) self.epe = EPE() self.sobel = SOBEL(backend) if local_rank != -1: self.flownet = DDP(self.flownet, device_ids=[local_rank], output_device=local_rank)

def eval(self): self.flownet.eval()

def device(self, backend): self.flownet.to(backend, non_blocking = True)

def inference(self, img0, img1, scale=1.0):
    imgs = torch.cat((img0, img1), 1)
    scale_list = [4/scale, 2/scale, 1/scale]
    _ , _ , merged = self.flownet(imgs, scale_list)
    return merged[2]



- In particular it gives this error on:
`mid_image = AI_model.inference(first_img, last_img)`

AI_model is an instance of the classe RIFEv3

Thank you

EDIT: it seems happening with other AI inference too :(

microsoft / DirectML

Check failed: rank <= DML_TENSOR_DIMENSION_COUNT_MAX #445