mirror of
https://github.com/wassname/Castor.git
synced 2026-09-10 11:40:44 +08:00
17 lines
389 B
Python
17 lines
389 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class LockedDropout(nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x, dropout=0.5):
|
|
if not self.training or not dropout:
|
|
return x
|
|
m = x.data.new(1, x.size(1), x.size(2)).bernoulli_(1 - dropout)
|
|
mask = m / (1 - dropout)
|
|
mask = mask.expand_as(x)
|
|
return mask * x
|