test1

contrastive_vae_tuner.py

@@ -27,16 +27,8 @@ import random
 import gc
 import argparse
 
-# Ensure PyTorch does NOT use GPUs
-os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # This disables CUDA for PyTorch
-torch_device = torch.device("cpu")  # Explicitly set PyTorch to CPU
 
-# Ensure TensorFlow CAN use GPUs
-os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"  # Optional: Prevent TF from using all memory
 tf_device = tf.config.list_physical_devices("GPU")  # Check available GPUs
-
-# Debugging Output
-print(f"PyTorch device: {torch_device}")
 print(f"TensorFlow devices: {tf_device}")  # Should list GPU(s)
 
 def distance(model, data):
@@ -60,7 +52,7 @@ def run(config):
     tf.config.experimental.enable_op_determinism()
 
     ####################################################################    
-    gpus = tf.config.experimental.list_physical_devices('CPU')
+    gpus = tf.config.experimental.list_physical_devices('GPU')
     if gpus:
         try:
             for gpu in gpus:
@@ -347,6 +339,6 @@ if __name__ == '__main__':
         search_alg=optuna_search,
         scheduler=scheduler,
         num_samples=1000,
-        resources_per_trial={'cpu': 8, 'gpu': 1 / 4},
+        resources_per_trial={'cpu': 2, 'gpu': 1 / 4},
     
     )