diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp
index 27394e18f9f14..e3544ca7e4858 100644
--- a/enzyme/Enzyme/EnzymeLogic.cpp
+++ b/enzyme/Enzyme/EnzymeLogic.cpp
@@ -2504,8 +2504,7 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
       Value *gep = ret;
       if (!removeStruct) {
 #if LLVM_VERSION_MAJOR > 7
-        gep = ib.CreateGEP(ret->getType()->getPointerElementType(), ret, Idxs,
-                           "");
+        gep = ib.CreateGEP(RetType, ret, Idxs, "");
 #else
         gep = ib.CreateGEP(ret, Idxs, "");
 #endif
@@ -2528,8 +2527,7 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
       tapeMemory = ret;
       if (!removeStruct) {
 #if LLVM_VERSION_MAJOR > 7
-        tapeMemory = ib.CreateGEP(ret->getType()->getPointerElementType(), ret,
-                                  Idxs, "");
+        tapeMemory = ib.CreateGEP(RetType, ret, Idxs, "");
 #else
         tapeMemory = ib.CreateGEP(ret, Idxs, "");
 #endif
@@ -2548,8 +2546,7 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
         Value *gep = tapeMemory;
         if (!removeTapeStruct) {
 #if LLVM_VERSION_MAJOR > 7
-          gep = ib.CreateGEP(tapeMemory->getType()->getPointerElementType(),
-                             tapeMemory, Idxs, "");
+          gep = ib.CreateGEP(tapeType, tapeMemory, Idxs, "");
 #else
           gep = ib.CreateGEP(tapeMemory, Idxs, "");
 #endif
@@ -2625,7 +2622,7 @@ const AugmentedReturn &EnzymeLogic::CreateAugmentedPrimal(
       ib.CreateRetVoid();
     else {
 #if LLVM_VERSION_MAJOR > 7
-      ib.CreateRet(ib.CreateLoad(ret->getType()->getPointerElementType(), ret));
+      ib.CreateRet(ib.CreateLoad(RetType, ret));
 #else
       ib.CreateRet(ib.CreateLoad(ret));
 #endif