How to use multithreading to speed up KL calculations？

[Horse4Truin]

I found that when I was doing KL calculation, only one thread was calculating, which is currently taking up a lot of my time. How can I solve this problem?

[Horse4Truin]

By the way，I don't know why the GPU setting is -1:ngpu=-1，here is my parameters
INFO:root:Writing training args to /home/guanaohan/foldingdiff/foldingdiff/results/training_args.json INFO:root:Training argument: results_dir=/home/guanaohan/foldingdiff/foldingdiff/results INFO:root:Training argument: dataset_key=cath INFO:root:Training argument: angles_definitions=canonical-full-angles INFO:root:Training argument: max_seq_len=128 INFO:root:Training argument: min_seq_len=40 INFO:root:Training argument: trim_strategy=randomcrop INFO:root:Training argument: timesteps=1000 INFO:root:Training argument: variance_schedule=cosine INFO:root:Training argument: variance_scale=1.0 INFO:root:Training argument: time_encoding=gaussian_fourier INFO:root:Training argument: num_hidden_layers=12 INFO:root:Training argument: hidden_size=384 INFO:root:Training argument: intermediate_size=768 INFO:root:Training argument: num_heads=12 INFO:root:Training argument: position_embedding_type=relative_key INFO:root:Training argument: dropout_p=0.1 INFO:root:Training argument: decoder=mlp INFO:root:Training argument: gradient_clip=1.0 INFO:root:Training argument: batch_size=64 INFO:root:Training argument: lr=5e-05 INFO:root:Training argument: loss=smooth_l1 INFO:root:Training argument: use_pdist_loss=0.0 INFO:root:Training argument: l2_norm=0.0 INFO:root:Training argument: l1_norm=0.0 INFO:root:Training argument: circle_reg=0.0 INFO:root:Training argument: min_epochs=10000 INFO:root:Training argument: max_epochs=10000 INFO:root:Training argument: early_stop_patience=0 INFO:root:Training argument: lr_scheduler=LinearWarmup INFO:root:Training argument: use_swa=False INFO:root:Training argument: subset=None INFO:root:Training argument: exhaustive_validation_t=False INFO:root:Training argument: syn_noiser= INFO:root:Training argument: single_angle_debug=-1 INFO:root:Training argument: single_timestep_debug=False INFO:root:Training argument: cpu_only=False INFO:root:Training argument: ngpu=-1 INFO:root:Training argument: write_valid_preds=False INFO:root:Training argument: dryrun=False INFO:root:Training argument: multithread=True
and my gpu setting is ok:
`>>> import torch

torch.cuda.is_available()
True
print(torch.cuda.device_count())
1
print(torch.cuda.get_device_name(0))
NVIDIA GeForce RTX 3080 Ti
print(torch.cuda.current_device())
0`