import torch
import torch.nn as nn
model = nn.Sequential(
nn.Linear(5,10),
nn.ReLU(),
nn.Dropout(p=0.5), # 50% chance of Dropout
nn.Linear(10,2),
nn.ReLU(),
nn.Dropout(p=0.2)
)
#call this on the model before training
#it activates the dropout layers
model.train()
#this will de-activate the dropout layers
#used after training to run in production
model.eval()