| 25 | |
| 26 | class Max_Embed(nn.Module): |
| 27 | def __init__(self, in_channels=2, out_channels=256, kernel_size = 3, stride = 1, padding = 1, shortcut= False): |
| 28 | super().__init__() |
| 29 | |
| 30 | self.embed_lif = MultiStepLIFNode(tau=2.0, detach_reset=True, backend='cupy') |
| 31 | self.embed_conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False) |
| 32 | self.embed_bn = nn.BatchNorm2d(out_channels) |
| 33 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False) |
| 34 | |
| 35 | self.shortcut = shortcut |
| 36 | |
| 37 | def forward(self, x): |
| 38 | #input : T, B, C, H, W |