def __init__(self, in_channels, channels, stride, erase_relu=False):
super().__init__()
out_channels = channels * self.expansion
self.conv1 = Conv2d(in_channels,
channels,
kernel_size=1,
norm='def',
act='def')
self.conv2 = Conv2d(channels,
channels,
kernel_size=3,
stride=stride,
norm='def',
act='def')
self.conv3 = Conv2d(channels, out_channels, kernel_size=1, norm='def')
if stride != 1 or in_channels != out_channels:
self.shortcut = Conv2d(in_channels,
out_channels,
kernel_size=1,
stride=stride,
norm='def')
else:
self.shortcut = Identity()
self.act = Act() if not erase_relu else Identity()
