diff --git a/ext/TensorOperationsChainRulesCoreExt.jl b/ext/TensorOperationsChainRulesCoreExt.jl
index eeee38a7..02bdae1b 100644
--- a/ext/TensorOperationsChainRulesCoreExt.jl
+++ b/ext/TensorOperationsChainRulesCoreExt.jl
@@ -27,12 +27,14 @@ using LinearAlgebra
 
 # Cannot free intermediate tensors when using AD
 # Thus we change the forward passes: `istemp=false` and `tensorfree!` is a no-op
-function ChainRulesCore.rrule(
-        ::typeof(TensorOperations.tensorfree!), allocator = DefaultAllocator()
-    )
+function ChainRulesCore.rrule(::typeof(TensorOperations.tensorfree!), C)
     tensorfree!_pullback(Δargs...) = (NoTangent(), NoTangent())
     return nothing, tensorfree!_pullback
 end
+function ChainRulesCore.rrule(::typeof(TensorOperations.tensorfree!), C, allocator)
+    tensorfree!_pullback(Δargs...) = (NoTangent(), NoTangent(), NoTangent())
+    return nothing, tensorfree!_pullback
+end
 function ChainRulesCore.rrule(
         ::typeof(TensorOperations.tensoralloc), ttype, structure,
         istemp, allocator = DefaultAllocator()
diff --git a/ext/TensorOperationsEnzymeExt/TensorOperationsEnzymeExt.jl b/ext/TensorOperationsEnzymeExt/TensorOperationsEnzymeExt.jl
index 8094b908..97c92ed6 100644
--- a/ext/TensorOperationsEnzymeExt/TensorOperationsEnzymeExt.jl
+++ b/ext/TensorOperationsEnzymeExt/TensorOperationsEnzymeExt.jl
@@ -9,6 +9,7 @@ using Enzyme.EnzymeCore
 using Enzyme.EnzymeCore: EnzymeRules
 
 @inline EnzymeRules.inactive(::typeof(TensorOperations.tensorfree!), ::Any) = true
+@inline EnzymeRules.inactive(::typeof(TensorOperations.tensorfree!), ::Any, ::Any) = true
 @inline EnzymeRules.inactive_type(v::Type{<:AbstractBackend}) = true
 @inline EnzymeRules.inactive_type(v::Type{DefaultAllocator}) = true
 @inline EnzymeRules.inactive_type(v::Type{<:CUDAAllocator}) = true
diff --git a/ext/TensorOperationsMooncakeExt/TensorOperationsMooncakeExt.jl b/ext/TensorOperationsMooncakeExt/TensorOperationsMooncakeExt.jl
index 57acd534..2ea5db53 100644
--- a/ext/TensorOperationsMooncakeExt/TensorOperationsMooncakeExt.jl
+++ b/ext/TensorOperationsMooncakeExt/TensorOperationsMooncakeExt.jl
@@ -27,6 +27,7 @@ Mooncake.tangent_type(::Type{ManualAllocator}) = Mooncake.NoTangent
 @zero_derivative DefaultCtx Tuple{typeof(TensorOperations.promote_add), Any}
 
 Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(TensorOperations.tensorfree!), Any}
+Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(TensorOperations.tensorfree!), Any, Any}
 Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(TensorOperations.tensoralloc), Any, Any, Any, Any}
 
 Mooncake.@is_primitive DefaultCtx Tuple{typeof(tensorcontract!), AbstractArray, AbstractArray, Index2Tuple, Bool, AbstractArray, Index2Tuple, Bool, Index2Tuple, Number, Number, Vararg{Any}}