set_platform

static Config.set_platform(platform)

Set the simulation platform of jax to 'cpu' or 'gpu'.

Parameters:

platform (str) – Simulation platform of jax.