From d39838bc7fec621c45fb156b2e62a74a513ba132 Mon Sep 17 00:00:00 2001
From: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
Date: Mon, 26 Sep 2022 13:40:05 -0400
Subject: [PATCH] skip zero/nan/inf checks on GPU

---
 src/layers/show.jl | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/layers/show.jl b/src/layers/show.jl
index 421131f365..ea852adc3a 100644
--- a/src/layers/show.jl
+++ b/src/layers/show.jl
@@ -113,7 +113,10 @@ underscorise(n::Integer) =
   join(reverse(join.(reverse.(Iterators.partition(digits(n), 3)))), '_')
 
 function _nan_show(io::IO, x)
-  if !isempty(x) && _all(iszero, x)
+  if any(y -> y isa Zygote.AbstractGPUArray, x)
+    # These friendly warnings take 10-20 sec to compile the first time, for models on GPU. 
+    printstyled(io, " (on GPU)", color=:light_black)
+  elseif !isempty(x) && _all(iszero, x)
     printstyled(io, "  (all zero)", color=:cyan)
   elseif _any(isnan, x)
     printstyled(io, "  (some NaN)", color=:red)