diff --git a/hpctestlib/microbenchmarks/gpu/gpu_burn.py b/hpctestlib/microbenchmarks/gpu/gpu_burn.py
index d1e1616398..3d677686aa 100644
--- a/hpctestlib/microbenchmarks/gpu/gpu_burn.py
+++ b/hpctestlib/microbenchmarks/gpu/gpu_burn.py
@@ -52,9 +52,9 @@ def setup_build(self):
 
         if self.gpu_build is None:
             # Try to set the build type from the partition features
-            if 'cuda' in curr_part.features:
+            if 'cuda' in curr_env.features:
                 self.gpu_build = 'cuda'
-            elif 'hip' in curr_part.features:
+            elif 'hip' in curr_env.features:
                 self.gpu_build = 'hip'
 
         gpu_devices = curr_part.select_devices('gpu')