Pass --gpus in packages/build.py

Author: nevillelyh

packages/build.py

@@ -37,6 +37,11 @@
     default='/tmp',
     help='Directory to mount read-write as /dst inside the container',
 )
+parser.add_argument(
+    '--gpus',
+    metavar='GPU',
+    help='docker run --gpus argument',
+)
 parser.add_argument(
     'script',
     metavar='SCRIPT',
@@ -117,6 +122,8 @@ def run(args: argparse.Namespace) -> None:
         '--volume',
         f'{os.path.abspath(args.dst)}:/dst:rw',
     ]
+    if args.gpus is not None:
+        run_cmd = run_cmd + ['--gpus', args.gpus]
     if args.script:
         run_cmd = run_cmd + [
             '--volume',

packages/torchsparse.sh

@@ -26,6 +26,10 @@ python3 setup.py bdist_wheel
 
 find . -name '*.whl' > requirements.txt
 # import torchspase fails without GPU device
-/build/test.sh requirements.txt 'print()' # 'import torchsparse.backend; "count_cuda" in torchsparse.backend.__dict__'
+if [ -e /dev/nvidiactl ]; then
+    /build/test.sh requirements.txt 'import torchsparse.backend; "count_cuda" in torchsparse.backend.__dict__'
+else
+    echo 'GPU not found, add "--gpus all"'
+fi
 
 find . -name '*.whl' -exec cp {} /dst \;