you can start the code with below command

# Nonlinear Diffusion Eq

XLA_PYTHON_CLIENT_PREALLOCATE=false CUDA_VISIBLE_DEVICES=0 python diffusion3d.py --data_dir=./data/diffusion3d --model=pinn --equation=diffusion3d --nc=16 --seed=100 --lr=0.001 --epochs=100000 --mlp=modified_mlp --n_layers=2 --features=16 --out_dim=1 --pos_enc=0 --log_iter=1000 --plot_iter=50000 --mlp_dim 4 --num_gaussian 4000 --grid_range 1. --sigmas_range 0.1