core.layer.supported_layers 

 1from torch import nn
 2
 3from core.layer import ProbBatchNorm1d, ProbBatchNorm2d, ProbConv2d, ProbLinear
 4
 5LAYER_MAPPING = {
 6    nn.Linear: ProbLinear,
 7    nn.Conv2d: ProbConv2d,
 8    nn.BatchNorm1d: ProbBatchNorm1d,
 9    nn.BatchNorm2d: ProbBatchNorm2d,
10}
LAYER_MAPPING = {<class 'torch.nn.modules.linear.Linear'>: <class 'core.layer.ProbLinear.ProbLinear'>, <class 'torch.nn.modules.conv.Conv2d'>: <class 'core.layer.ProbConv2d.ProbConv2d'>, <class 'torch.nn.modules.batchnorm.BatchNorm1d'>: <class 'core.layer.ProbBatchNorm1d.ProbBatchNorm1d'>, <class 'torch.nn.modules.batchnorm.BatchNorm2d'>: <class 'core.layer.ProbBatchNorm2d.ProbBatchNorm2d'>}