Fix package version logic (pytorch#2048)

Summary:
Pull Request resolved: pytorch#2048

- Fix package version logic
- Change default package variant to cuda
- Fix rocm variants for testing (rocm versions on Nova are 5.6 and 5.7)

Reviewed By: sryap

Differential Revision: D49740927

fbshipit-source-id: 870812fd46dbd3e5e2765540b94fb45951811a8f

.github/workflows/fbgemm_gpu_pip.yml

@@ -172,7 +172,7 @@ jobs:
         ]
         # ROCm machines are limited, so we only test against Python 3.10
         python-version: [ "3.10" ]
-        rocm-version: [ "5.5.1", "5.6" ]
+        rocm-version: [ "5.6", "5.7" ]
 
     steps:
     - name: Setup Build Container

fbgemm_gpu/setup.py

@@ -30,7 +30,7 @@ def parse_args(argv: List[str]) -> argparse.Namespace:
         "--package_variant",
         type=str,
         choices=["cpu", "cuda", "rocm"],
-        default="cpu",
+        default="cuda",
         help="The FBGEMM_GPU variant to build.",
     )
     parser.add_argument(
@@ -344,7 +344,7 @@ def main(argv: List[str]) -> None:
 
     # Generate the full package version string
     package_version = FbgemmGpuInstaller.generate_package_version(
-        package_name, variant_version
+        args.package_name, variant_version
     )
 
     # Generate the version file