Light-amplification

Running on Zero

Fediory commited on 27 days ago

Commit

0f856ef

1 Parent(s): ab2b5fe

fix: unable find cuda

Files changed (1) hide show

net/CIDNet.py CHANGED Viewed

@@ -67,7 +67,7 @@ class CIDNet(nn.Module):
         self.I_LCA5 = I_LCA(ch3, head3)
         self.I_LCA6 = I_LCA(ch2, head2)
-        self.trans = RGB_HVI().cuda()
     def forward(self, x):
         dtypes = x.dtype

         self.I_LCA5 = I_LCA(ch3, head3)
         self.I_LCA6 = I_LCA(ch2, head2)
+        self.trans = RGB_HVI()
     def forward(self, x):
         dtypes = x.dtype