From 7ab9d8ed3d3609c0a42364ccaa8ba95fa4df27de Mon Sep 17 00:00:00 2001
From: thebhatman <manjunathbhat9920@gmail.com>
Date: Thu, 13 Jun 2019 18:59:03 +0530
Subject: [PATCH] Minor update

---
 src/layers/normalise.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/layers/normalise.jl b/src/layers/normalise.jl
index 1adc3050..3755f3fc 100644
--- a/src/layers/normalise.jl
+++ b/src/layers/normalise.jl
@@ -29,7 +29,7 @@ _dropout_kernel(y::T, p, q) where {T} = y > p ? T(1 / q) : T(0)
 
 function dropout(x, p; dims = :)
   istraining() || return x
-  y = similar(x)
+  y = similar(x, _dropout_shape(x, dims))
   rand!(y)
   y .= _dropout_kernel.(y, p, 1 - p)
   return x .* y