CUDA issues

[RaresAmbrus]

Thank you for your work and for open-sourcing the code! I am trying to run your pipeline, but I am running into a couple of issues. I have followed the instructions from the repo to set up the code as well as the jittor instructions to enable CUDA support.

First, it seems like train_part.py isn't using CUDA, and it's running very slowly on my machine (1% in ~8 hours). If I add jt.flags.use_cuda = 1 I get the follwoing error:

python3.7: symbol lookup error: /home/ubuntu/.cache/jittor/jt1.3.4/g++7.5.0/py3.7.13/Linux-5.4.0-10xfe/IntelRCoreTMi9x9f/default/cu11.2.152_sm_86/jit/cudnn_conv3d_Tx_bool__Ty_float32__Tw_float32__JIT_1__JIT_cuda_1__index_t_int32__hash_b2c2838050d33db1_op.so: undefined symbol: _ZN6jittor11getDataTypeIbEE15cudnnDataType_tv

On the other hand, I can run python3.7 -m jittor.test.test_cudnn_op and python3.7 -m jittor.test.test_resnet without any errors.

Any ideas? Thanks for the help!