From 29230507d7d1b46d45c5d14eac86f03f903a274c Mon Sep 17 00:00:00 2001
From: Lukas Billera <125421968+billera@users.noreply.github.com>
Date: Tue, 31 Dec 2024 05:57:35 +0100
Subject: [PATCH] add arg

---
 src/layers.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/layers.jl b/src/layers.jl
index 5e8753a..38de2f6 100644
--- a/src/layers.jl
+++ b/src/layers.jl
@@ -268,7 +268,7 @@ function (ipa::Union{IPCrossA, IPA})(
     return si 
 end
 
-function ipa_customgrad(ipa::Union{IPCrossA, IPA}, Ti::Tuple{AbstractArray,AbstractArray}, S::AbstractArray, zij::AbstractArray, mask::AbstractArray, rope::Union{IPARoPE, Nothing} = nothing)    
+function ipa_customgrad(ipa::Union{IPCrossA, IPA}, Ti::Tuple{AbstractArray,AbstractArray}, S::AbstractArray, zij::AbstractArray, mask::AbstractArray, rope::Union{IPARoPE, Nothing} = nothing, chain_diffs = 1)    
     # Get relevant parameters from our ipa struct.
     l = ipa.layers
     dims, c, N_head, N_query_points, N_point_values, c_z, Typ, pairwise = ipa.settings