Update WT/transform.py

WT/transform.py

@@ -23,7 +23,7 @@ def iwt_init(x):
     x2 = x[:, out_channel:out_channel * 2, :, :] / 2
     x3 = x[:, out_channel * 2:out_channel * 3, :, :] / 2
     x4 = x[:, out_channel * 3:out_channel * 4, :, :] / 2
-    h = torch.zeros([out_batch, out_channel, out_height, out_width]).cuda() #
+    h = torch.zeros([out_batch, out_channel, out_height, out_width])#.cuda() #
 
     h[:, :, 0::2, 0::2] = x1 - x2 - x3 + x4
     h[:, :, 1::2, 0::2] = x1 - x2 + x3 - x4