From 02a463832e16423849ad5a5f55bf5abd7972fc54 Mon Sep 17 00:00:00 2001
From: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
Date: Fri, 14 Oct 2022 22:30:25 -0400
Subject: [PATCH] fix |> gpu bug in autosize

---
 src/outputsize.jl  | 7 ++++---
 test/outputsize.jl | 7 +++++--
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/outputsize.jl b/src/outputsize.jl
index f952a223c2..76c58237be 100644
--- a/src/outputsize.jl
+++ b/src/outputsize.jl
@@ -204,7 +204,7 @@ julia> @autosize (img..., 1, 32) Chain(              # size is only needed at ru
           Dense(_ => _÷4, relu, init=Flux.rand32),   # can calculate output size _÷4
           SkipConnection(Dense(_ => _, relu), +),
           Dense(_ => 10),
-       ) |> gpu                                      # moves to GPU after initialisation
+       )
 Chain(
   Chain(
     c = Conv((3, 3), 1 => 5, pad=1, stride=2),  # 50 parameters
@@ -290,8 +290,6 @@ mutable struct LazyLayer
   layer
 end
 
-@functor LazyLayer
-
 function (l::LazyLayer)(x::AbstractArray, ys::AbstractArray...)
   l.layer === nothing || return l.layer(x, ys...)
   made = l.make(x)  # for something like `Bilinear((_,__) => 7)`, perhaps need `make(xy...)`, later.
@@ -320,6 +318,9 @@ function ChainRulesCore.rrule(::typeof(striplazy), m)
 end
 
 params!(p::Params, x::LazyLayer, seen = IdSet()) = error("LazyLayer should never be used within params(m). Call striplazy(m) first.")
+
+Functors.functor(::Type{<:LazyLayer}, x) = error("LazyLayer should not be walked with Functors.jl, as the arrays which Flux.gpu wants to move may not exist yet.")
+
 function Base.show(io::IO, l::LazyLayer)
   printstyled(io, "LazyLayer(", color=:light_black)
   if l.layer == nothing
diff --git a/test/outputsize.jl b/test/outputsize.jl
index 2e9595f699..eec6880dc2 100644
--- a/test/outputsize.jl
+++ b/test/outputsize.jl
@@ -220,8 +220,8 @@ end
          Dense(_ => _÷4, relu, init=Flux.rand32),   # can calculate output size _÷4
          SkipConnection(Dense(_ => _, relu), +),
          Dense(_ => 10),
-      ) |> gpu                                      # moves to GPU after initialisation
-  @test randn(Float32, img..., 1, 32) |> gpu |> m |> size == (10, 32)
+      )
+  @test randn(Float32, img..., 1, 32) |> m |> size == (10, 32)
 end
 
 @testset "LazyLayer" begin
@@ -241,4 +241,7 @@ end
   @test_throws Exception Flux.params(lm)
   @test_throws Exception gradient(x -> sum(abs2, lm(x)), [1,2])
   @test_throws Exception gradient(m -> sum(abs2, Flux.striplazy(m)([1,2])), ld)
+  
+  # Can't let |> gpu act before the arrays are materialized... so it's an error: 
+  @test_throws ErrorException @eval @autosize (1,2,3) Dense(_=>2) |> f64
 end