Spaces:
Sleeping
Sleeping
cuda_condition initialized
Browse files
new_test_saved_finetuned_model.py
CHANGED
@@ -46,7 +46,7 @@ class BERTFineTuneTrainer:
|
|
46 |
"""
|
47 |
|
48 |
# Setup cuda device for BERT training, argument -c, --cuda should be true
|
49 |
-
|
50 |
self.device = torch.device("cuda:0" if cuda_condition else "cpu")
|
51 |
#self.device = torch.device("cpu") #torch.device("cuda:0" if cuda_condition else "cpu")
|
52 |
# print(cuda_condition, " Device used = ", self.device)
|
|
|
46 |
"""
|
47 |
|
48 |
# Setup cuda device for BERT training, argument -c, --cuda should be true
|
49 |
+
cuda_condition = torch.cuda.is_available() and with_cuda
|
50 |
self.device = torch.device("cuda:0" if cuda_condition else "cpu")
|
51 |
#self.device = torch.device("cpu") #torch.device("cuda:0" if cuda_condition else "cpu")
|
52 |
# print(cuda_condition, " Device used = ", self.device)
|