@@ -184,25 +184,25 @@ for runtime in (:PJRT, :IFRT)
184
184
state. default_client = cpu
185
185
186
186
# Try TPU if possible, then try GPU (CUDA)
187
- if ! Reactant. precompiling ()
188
- @static if ! Sys. isapple ()
189
- if Reactant. has_tpu ()
190
- dataset_dir = @get_scratch! (" libtpu" )
191
- download_tpu (dataset_dir)
192
- try
193
- if was_initialized && haskey (state. clients, " tpu" )
194
- XLA. free_client (state. clients[" tpu" ])
195
- XLA.$ (runtime). tpu_client_count[] -= 1
187
+ if ! Reactant. precompiling ()
188
+ @static if ! Sys. isapple ()
189
+ if Reactant. has_tpu ()
190
+ dataset_dir = @get_scratch! (" libtpu" )
191
+ download_tpu (dataset_dir)
192
+ try
193
+ if was_initialized && haskey (state. clients, " tpu" )
194
+ XLA. free_client (state. clients[" tpu" ])
195
+ XLA.$ (runtime). tpu_client_count[] -= 1
196
+ end
197
+ tpu = $ (runtime). TPUClient (;
198
+ tpu_path= dataset_dir * " /libtpu.so" , common_kwargs...
199
+ )
200
+ state. clients[" tpu" ] = tpu
201
+ state. default_client = tpu
202
+ catch e
203
+ println (stdout , e)
196
204
end
197
- tpu = $ (runtime). TPUClient (;
198
- tpu_path= dataset_dir * " /libtpu.so" , common_kwargs...
199
- )
200
- state. clients[" tpu" ] = tpu
201
- state. default_client = tpu
202
- catch e
203
- println (stdout , e)
204
- end
205
- else
205
+ else
206
206
try
207
207
if was_initialized && haskey (state. clients, " gpu" )
208
208
XLA. free_client (state. clients[" gpu" ])
0 commit comments