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()