fixed : tf call cuda

This commit is contained in:
Zchen
2025-10-15 23:37:24 +08:00
parent 01024678c1
commit f9d3f47d20
2 changed files with 149 additions and 1 deletions

View File

@@ -763,6 +763,14 @@ def create_tpu_strategy():
print("🔍 Detecting TPU environment...")
# Disable GPU to avoid CUDA conflicts in TPU environment
try:
print("🚫 Disabling GPU to prevent CUDA conflicts...")
tf.config.set_visible_devices([], 'GPU')
print("✅ GPU disabled successfully")
except Exception as e:
print(f"⚠️ Warning: Could not disable GPU: {e}")
# Check for various TPU environment variables
tpu_address = None
tpu_name = None