Hey Apil, you can always run the command nvidia-smi to get an idea of what’s going on with your GPU.
I think we had similar specs (1080 Ti) and I started hitting similar issues when I ran that particular fit function after running all the previous cells in the notebook. It might be a memory problem. I was able to get it to work by restarting the kernel and running the first 5 commands (all the import commands and initializing global vars). Then skipping down and running these next: