Merge pull request #75 from slimgroup/mix-type-rr

Mix type rr

Author: mloubout

Project.toml

@@ -1,7 +1,7 @@
 name = "InvertibleNetworks"
 uuid = "b7115f24-5f92-4794-81e8-23b0ddb121d3"
 authors = ["Philipp Witte <[email protected]>", "Ali Siahkoohi <[email protected]>", "Mathias Louboutin <[email protected]>", "Gabrio Rizzuti <[email protected]>", "Rafael Orozco <[email protected]>", "Felix J. herrmann <[email protected]>"]
-version = "2.2.3"
+version = "2.2.4"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

src/utils/chainrules.jl

@@ -70,7 +70,7 @@ function forward_update!(state::InvertibleOperationsTape, X::AbstractArray{T,N},
 
     if isa_newblock(state, X)
         push!(state.Y, Y)
-        push!(state.layer_blocks, [net])
+        push!(state.layer_blocks, Vector{Any}([net]))
         state.counter_block += 1
         state.counter_layer = 1
     else

src/utils/compute_utils.jl

@@ -21,7 +21,7 @@ function chain_lr(x::AbstractMatrix{T}, vi::Vararg{AbstractVector{T}, N}) where
     out = T(1) .* x
     tmp = cuzeros(vi[1], size(x, 1))
     for v=vi
-        n = -2/norm(v)^2
+        n = -2/dot(v, v)
         mul!(tmp, out, v)
         rmul!(tmp, n)
         gemm_outer!(out, tmp, v)