fix device

Author: hanjr92

tensorlayerx/backend/ops/torch_nn.py

@@ -2182,7 +2182,7 @@ def __init__(self, data_format):
         self.data_format = data_format
 
     def __call__(self, input, weight):
-        weight = weight.to(input.device)
+        # weight = weight.to(input.device)
         return torch.prelu(input, weight)
 
 

tensorlayerx/dataflow/utils.py

@@ -78,22 +78,22 @@ def default_collate_torch(batch):
     data = batch[0]
     data_type = type(data)
     import torch
-    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    # device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
     if isinstance(data, torch.Tensor):
         batch = torch.stack(batch, 0)
-        batch = batch.to(device)
+        # batch = batch.to(device)
         return batch
     elif isinstance(data, np.ndarray):
         batch = np.stack(batch, axis=0)
         batch = torch.as_tensor(batch)
-        batch = batch.to(device)
+        # batch = batch.to(device)
         return batch
     elif isinstance(data, numbers.Number):
         batch = torch.as_tensor(batch)
-        batch = batch.to(device)
+        # batch = batch.to(device)
         return batch
     elif isinstance(data, (str, bytes)):
-        batch = batch.to(device)
+        # batch = batch.to(device)
         return batch
     elif isinstance(data, collections.abc.Mapping):
         return {key: default_collate_torch([d[key] for d in batch]) for key in data}

tensorlayerx/model/core.py

@@ -429,8 +429,8 @@ def th_train(
         self, n_epoch, train_dataset, network, loss_fn, train_weights, optimizer, metrics, print_train_batch,
         print_freq, test_dataset
     ):
-        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-        network = network.to(device)
+        # device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        # network = network.to(device)
         for epoch in range(n_epoch):
             start_time = time.time()
 

tensorlayerx/model/utils.py

@@ -177,20 +177,20 @@ def __call__(self, data, label):
 class TrainOneStepWithTH(object):
 
     def __init__(self, net_with_loss, optimizer, train_weights):
-        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        # self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
         self.net_with_loss = net_with_loss
-        self.net_with_loss = self.net_with_loss.to(self.device)
+        # self.net_with_loss = self.net_with_loss.to(self.device)
         self.optimizer = optimizer
         self.train_weights = train_weights
 
     def __call__(self, data, label):
-        if isinstance(data, dict):
-            for k, v in data.items():
-                if isinstance(v, torch.Tensor):
-                    data[k] = v.to(self.device)
-        elif isinstance(data, torch.Tensor):
-            data = data.to(self.device)
-        label = label.to(self.device)
+        # if isinstance(data, dict):
+        #     for k, v in data.items():
+        #         if isinstance(v, torch.Tensor):
+        #             data[k] = v.to(self.device)
+        # elif isinstance(data, torch.Tensor):
+        #     data = data.to(self.device)
+        # label = label.to(self.device)
         loss = self.net_with_loss(data, label)
         grads = self.optimizer.gradient(loss, self.train_weights)
         self.optimizer.apply_gradients(zip(grads, self.train_weights))