diff --git a/.github/workflows/build-push-base.yaml b/.github/workflows/build-push-base.yaml
index 4eb5831..62ffd86 100644
--- a/.github/workflows/build-push-base.yaml
+++ b/.github/workflows/build-push-base.yaml
@@ -63,23 +63,14 @@ jobs:
       matrix:
         target:
           - base
+          - trt
         cuda:
           - cu118
           - cu121
         torch:
           - torch201
           - torch210
-        exclude:
-          - target: base
-            cuda: cu121
-            torch: torch201
         include:
-          - target: trt
-            cuda: cu121
-            torch: torch201
-          - target: trt
-            cuda: cu121
-            torch: torch210
           - target: base
             cuda: cu121
             torch: nightly
diff --git a/docker-bake.hcl b/docker-bake.hcl
index c55f6ff..7c27c4f 100644
--- a/docker-bake.hcl
+++ b/docker-bake.hcl
@@ -150,7 +150,7 @@ target "base" {
     CUDA_RELEASE = cudaRelease(cuda.version)
 
     TORCH_INDEX      = "${torch.index}/${cudaName(cuda.version)}"
-    TORCH_PACKAGE    = "torch==${torch.version}+${cudaName(cuda.version)}"
+    TORCH_PACKAGE    = "torch"
     TRITON_PACKAGE   = torch.triton
     XFORMERS_PACKAGE = torch.xformers
     INCLUDE_TRT      = cuda.with-trt