Fix index access to new OneHotArray

FluxML · Jun 15, 2021 · e42a0d5 · e42a0d5
1 parent a52d343
commit e42a0d5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/onehot.jl b/src/onehot.jl
@@ -193,14 +193,14 @@ for wrapper in [:Adjoint, :Transpose]
       size(A, 2) == L ||
         throw(DimensionMismatch("Matrix column must correspond with OneHot size: $(size(A, 2)) != $L"))
 
-      return A[:, b.ix]
+      return A[:, onecold(b)]
     end
 
     function Base.:*(A::$wrapper{<:Number, <:AbstractVector{T}}, b::OneHotVector{<:Any, L}) where {L, T}
       size(A, 2) == L ||
         throw(DimensionMismatch("Matrix column must correspond with OneHot size: $(size(A, 2)) != $L"))
 
-      return A[b.ix]
+      return A[onecold(b)]
     end
   end
 end