From 77309609faec5e45ed338b4d91d475b0449fd128 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Wed, 11 Aug 2021 14:11:44 +0800
Subject: [PATCH] perf(functional/dropout): add fastpath for dropout

GitOrigin-RevId: 3bf8546908c2cd41a9d33c6236107a15e53f9fb4
---
 imperative/python/megengine/functional/nn.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/imperative/python/megengine/functional/nn.py b/imperative/python/megengine/functional/nn.py
index b87ded3e..72078a7b 100644
--- a/imperative/python/megengine/functional/nn.py
+++ b/imperative/python/megengine/functional/nn.py
@@ -1304,6 +1304,8 @@ def dropout(inp: Tensor, drop_prob: float, training: bool = True) -> Tensor:
 
     """
     assert 0 <= drop_prob < 1
+    if drop_prob == 0:
+        return inp
     rv = uniform(size=inp.shape)
     mask = rv > drop_prob
     inp *= mask.astype(inp.dtype)