VAE can't train

[LeeEuShane]

Describe the bug
train_vae.py is not working. I'm using the images which I generated using generate_data.py.
To Reproduce
Steps to reproduce the behavior:
train_vae.py --dir dataset/CartPoleContinuousBulletEnv-v0

Expected behavior
I'm not very sure what is expected to happen as I have not succeeded yet in making it work.

Screenshots
If applicable, add screenshots to help explain your problem.

Desktop (please complete the following information):

• OS: Linux (Workstation) ; Windows 10 (PC)
• Browser: Chrome
• Version: 93.0.4577.63

Smartphone (please complete the following information):

• Device: iphone 11
• OS: iOS 14.71
• Browser: Safari
• Version: Can't find but it's the latest

Additional context:
My understanding is to use the generated pictures and encode it to smaller dimension size to save time during the training. If my understanding is wrong, I may have used the train_vae.py script wrongly.