Has anyone had this issue with batch GPU

OSError: libcufft.so.11: cannot open shared object file: No such file or directory

and 

ValueError: libcublas.so.[0-9] not found in the system path ['/metaflow', '/metaflow/_escape_trampolines',

Built an environment using the pypi extension. Batch correctly spins up a

g4dn.2xlarge

and things start running. I figure I'm not passing a env variable or something to get cuda installed. Full text in the thread