From 60cd4f63c55c77f5097ebac146c212763ba11925 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Fran=C3=A7ois=20Fleuret?= Date: Fri, 19 Jan 2024 14:11:45 +0100 Subject: [PATCH] Update. --- mygpt.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/mygpt.py b/mygpt.py index 0414bb6..760a3c6 100755 --- a/mygpt.py +++ b/mygpt.py @@ -629,15 +629,21 @@ class Caterpillar(nn.Module): warnings.warn("gate dropout", RuntimeWarning) + if self.gate_dropout_sync: + shape_kill = (N, 1, 1) + else: + shape_kill = (N, H, R) + # Pick a point in each of the NxHxR timeline and set this # entry and the following to 1 kill = ( - torch.rand(N, H, R, t1 - t0, device=G.device).sort(dim=3).indices == 0 + torch.rand(*shape_kill, t1 - t0, device=G.device).sort(dim=3).indices + == 0 ).cumsum(dim=3) # Keep these mask for only some of the NxHxR kill = kill * ( - torch.rand(N, H, R, 1, device=G.device) <= self.gate_dropout_proba + torch.rand(*shape_kill, 1, device=G.device) <= self.gate_dropout_proba ) # The coefficient to keep are the complementary -- 2.39.5