@ptrblck Thanks for the response. As you had correctly mentioned, the issue with that case was environment-related and was caused by conflicting libraries.
Last night I tried running the code in this question again, and it ran smoothly. Since when I log into the remote machine, I get assigned to a randomly selected GPU, perhaps the issue was from the particular GPU node that I was assigned to in that session.