diff --git a/drivers/gpu/nvgpu/hal/class/class_gv11b_fusa.c b/drivers/gpu/nvgpu/hal/class/class_gv11b_fusa.c
index 379c554ad..3661143a8 100644
--- a/drivers/gpu/nvgpu/hal/class/class_gv11b_fusa.c
+++ b/drivers/gpu/nvgpu/hal/class/class_gv11b_fusa.c
@@ -33,12 +33,16 @@ bool gv11b_class_is_valid(u32 class_num)
 	nvgpu_speculation_barrier();
 
 	switch (class_num) {
-	case VOLTA_A:
 	case VOLTA_COMPUTE_A:
 	case VOLTA_DMA_COPY_A:
 	case VOLTA_CHANNEL_GPFIFO_A:
 		valid = true;
 		break;
+#ifdef CONFIG_NVGPU_GRAPHICS
+	case VOLTA_A:
+		valid = true;
+		break;
+#endif
 	default:
 		valid = gp10b_class_is_valid(class_num);
 		break;
diff --git a/drivers/gpu/nvgpu/include/nvgpu/class.h b/drivers/gpu/nvgpu/include/nvgpu/class.h
index bf31d6d07..951a1dfdc 100644
--- a/drivers/gpu/nvgpu/include/nvgpu/class.h
+++ b/drivers/gpu/nvgpu/include/nvgpu/class.h
@@ -69,16 +69,6 @@
  */
 #define	PASCAL_DMA_COPY_A		0xC0B5U
 
-/**
- * @ingroup NVGPU_CLASS_VALID_NUM
- *
- * Class number for 3D graphics class methods on Volta chips.
- *
- * WAR: Lot of qnx safety tests are still using graphics 3d class. Until these
- *      tests get fixed, allowing 3d graphics class as valid class for
- *      safety build.
- */
-#define	VOLTA_A				0xC397U
 /**
  * @ingroup NVGPU_CLASS_VALID_NUM
  *
@@ -102,6 +92,7 @@
 #define	FERMI_TWOD_A			0x902DU
 #define	MAXWELL_B			0xB197U
 #define	PASCAL_A			0xC097U
+#define	VOLTA_A				0xC397U
 #define	TURING_A			0xC597U
 #endif