How to run with accelerate using TPU ?

[Pham-Xuan-Thanh]

Im get error when train_sde.py when it run function import_modules_into_registry in advanced of train. That check cuda version while im running tpu. So i got an error:

CUDA_HOME environment variable is not set. Please set it to your CUDA install root. File "/Users/user/Desktop/master/CV/PSLD/main/models/score_fn/song_sde/op/fused_act.py", line 11, in <module> fused = load( File "/Users/user/Desktop/master/CV/PSLD/main/models/score_fn/song_sde/op/__init__.py", line 1, in <module> from .fused_act import FusedLeakyReLU, fused_leaky_relu File "/Users/user/Desktop/master/CV/PSLD/main/models/score_fn/song_sde/up_or_down_sampling.py", line 10, in <module> from .op import upfirdn2d File "/Users/user/Desktop/master/CV/PSLD/main/models/score_fn/song_sde/layerspp.py", line 20, in <module> from . import up_or_down_sampling File "/Users/user/Desktop/master/CV/PSLD/main/models/score_fn/song_sde/ncsnpp.py", line 18, in <module> from . import utils, layers, layerspp, normalization File "/Users/user/Desktop/master/CV/PSLD/main/models/__init__.py", line 2, in <module> from .score_fn.song_sde.ncsnpp import NCSNpp File "/Users/user/Desktop/master/CV/PSLD/main/util.py", line 120, in import_modules_into_registry import models File "/Users/user/Desktop/master/CV/PSLD/main/train_sde.py", line 18, in <module> import_modules_into_registry() OSError: CUDA_HOME environment variable is not set. Please set it to your CUDA install root.