getting error on line line 148 in main.py

[M-Suzy]

Hi,
Could you please help to run your code without errors.
Right now I have the following issue:
2021-02-23 21:25:34,587 INFO Loading Adobe5k dataset ... Traceback (most recent call last): File "main.py", line 351, in <module> main() File "main.py", line 148, in main inference_dataset = Dataset(data_dict=inference_data_dict, NameError: name 'Dataset' is not defined

[sjmoran]

Hi, thank you for your interest in our work. I will have updated code in a few days time, including a new pre-trained model.

[M-Suzy]

Yes, I am using the updated commit. Here is the command I use which gives the error
python3.8 main.py --inference_img_dirpath=./adobe5k_dpe/ --checkpoint_filepath=./pretrained_models/curl_validpsnr_23.070664790618277_validloss_0.029170328751206398_testpsnr_23.556465119481636_testloss_0.02679653838276863_epoch_24_model.pt

[sjmoran]

Issue now resolved, please see the latest code pushed to the repository today. I have also uploaded a model from epoch 99 that is more accurate (better PSNR, SSIM).

[M-Suzy]

Thank you very much for the fix! By the way, does your code support only GPU, or I can also run on CPU mode?

[sjmoran]

Hi, CPU mode is straightforward. You just need to place the cuda keywords in the code with cpu I believe e.g. net.cuda() with net.cpu() etc. For inference you'll also need to load the model on the CPU, more information is here: https://pytorch.org/tutorials/beginner/saving_loading_models.html. Feel free to raise a PR if you have a CPU/GPU switch implemented for CURL.

[M-Suzy]

Got it, thank you!