Leveraging a GPU in Python Compute Modules

We are trying to use GPU libraries for pytorch and cuda on a python compute module. We are trying to detect GPU access in the compute module, but the function to detect ability to use GPU via pytorch is returning false. When I check the lockfile of the repo, it seems like we are running into this error where the CPU-only packages are installed. However, the suggested fix in a code workspace is to set the job to run with a GPU directly in the repo, which makes sense for something like a Jupyter notebook. However, with a compute module, you set up the GPU after the container has already been built. Is there a way to force my compute module python repo to compile with the GPU-relevant libraries so that I can leverage the GPU inside my compute module? Any examples for how to use the GPU from within a python compute module?