Update depth_decoder.py

depth_decoder.py

@@ -18,6 +18,8 @@ class DepthDecoder(nn.Module):
     def __init__(self, num_ch_enc, scales=range(4), num_output_channels=1, use_skips=True, batch_norm = True):
         super(DepthDecoder, self).__init__()
 
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
         self.num_output_channels = num_output_channels
         self.use_skips = use_skips
         self.upsample_mode = 'nearest'
@@ -70,7 +72,7 @@ class DepthDecoder(nn.Module):
             x = torch.cat(x, 1)
             x = self.convs[("upconv", i, 1)](x)
             if self.batch_norm:
-                x = self.bn[('bn', i)].cuda()(x)
+                x = self.bn[('bn', i)].to(self.device)(x)
             
                 
             # batchnorm