Update batching.jl

Author: murrellb

src/batching.jl

@@ -1,13 +1,13 @@
+#This will be in Flowfusion:
 element(state, seqindex) = selectdim(state, ndims(state), seqindex:seqindex)
 element(state, seqindex, batchindex) = element(selectdim(state, ndims(state), batchindex), seqindex)
-
 element(S::MaskedState, inds...) = element(S.S, inds...)
 element(S::ContinuousState, inds...) = ContinuousState(element(S.state, inds...))
 element(S::ManifoldState, inds...) = ManifoldState(S.M, element(S.state, inds...))
 element(S::DiscreteState, inds...) = DiscreteState(S.K, element(S.state, inds...))
-
 element(S::Tuple{Vararg{Flowfusion.UState}}, inds...) = element.(S, inds...)
 
+
 #Create a "zero" state appropriate for the type. Tricky for manifolds, but we just want rotations working for now I think.
 zerostate(element::T, expandsize...) where T <: ContinuousState = T(similar(tensor(element), size(tensor(element))[1:end-1]..., expandsize...) .= 0)
 zerostate(element::DiscreteState{<:AbstractArray{<:Signed}}, expandsize...) = DiscreteState(element.K,similar(tensor(element), size(tensor(element))[1:end-1]..., expandsize...) .= element.K)
@@ -31,11 +31,7 @@ function regroup(elarray::AbstractArray{<:AbstractArray})
     example_tuple = elarray[1][1]
     maxlen = maximum(length.(elarray))
     b = length(elarray)
-    @show maxlen, b
     newstates = [zerostate(example_tuple[i],maxlen,b) for i in 1:length(example_tuple)]
-    for ne in newstates
-        @show size(tensor(ne))
-    end
     for i in 1:b
         for j in 1:length(elarray[i])
             for k in 1:length(example_tuple)
@@ -44,4 +40,4 @@ function regroup(elarray::AbstractArray{<:AbstractArray})
         end
     end
     return Tuple(newstates)
-end
+end