Canny kornia RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.HalfTensor) should be the same

I’m trying to modify a module in YoloV9, but when I use the Canny function from Kornia. I got this error


Here is my code :

class SpectralFeatureAdaptation(nn.Module):
    def __init__(self,c1):
        super().__init__()
        self.conv1= nn.Conv3d(4, 16, 3, stride=1, padding='same')  ##in=4 if edge info used
        self.conv2= nn.Conv3d(4, 16, (3,5,5), stride=1, padding='same') ##in=4 if edge info used
        self.conv3= nn.Conv3d(4, 16, (3,7,7), stride=1, padding='same') ##in=4 if edge info used
        self.conv4= nn.Conv3d(48, 3, 3, stride=1, padding='same')
   
    def forward(self, input):
        mag, edge = canny(input)
        edge_fused=torch.cat((input,edge), 1)
        ba,ch,dim1,dim2=list(edge_fused.size())
        input=input.reshape(ba,ch,-1,dim1,dim2)
        x1= self.conv1(edge_fused)
        x2= self.conv2(edge_fused)
        x3= self.conv3(edge_fused)
        x4=torch.cat((x1,x2,x3), 1) ## concatinate feature map
        x5= self.conv4(x4)
        ba,ch,di,dim1,dim2=list(x5.size())
        x5=x5.reshape(ba,ch,dim1,dim2)
        return x5

Based on the error message it seems canny contains float16 parameters while input is in float32 raising the dtype mismatch. You could add a torch.cuda.amp.autocast context to this operation in case you want to use mixed-precision execution.

Also, it’s better to post code snippets instead of screenshots as they are not only hard to read but also cannot be indexed into our search. canny is btw. also not defined, so I don’t know how its dtype is defined.