(self, feat_maps)
| 235 | self.shortcut = nn.Identity() |
| 236 | |
| 237 | def forward(self, feat_maps): |
| 238 | if self.input_norm: |
| 239 | feat_maps = [self.norm_xy(feat_maps[0]), self.norm_xz(feat_maps[1]), self.norm_yz(feat_maps[2])] |
| 240 | x, (H, W, D) = compose_triplane_channelwise(feat_maps) |
| 241 | |
| 242 | if self.up: |
| 243 | raise NotImplementedError |
| 244 | else: |
| 245 | h = self.in_layers(x) |
| 246 | |
| 247 | h_xy, h_xz, h_yz = decompose_triplane_channelwise(h, (H, W, D)) |
| 248 | h_xy = self.norm_xy(h_xy) |
| 249 | h_xz = self.norm_xz(h_xz) |
| 250 | h_yz = self.norm_yz(h_yz) |
| 251 | h, _ = compose_triplane_channelwise([h_xy, h_xz, h_yz]) |
| 252 | |
| 253 | h = self.out_layers(h) |
| 254 | h = h + self.shortcut(x) |
| 255 | h_maps = decompose_triplane_channelwise(h, (H, W, D)) |
| 256 | return h_maps |
nothing calls this directly
no test coverage detected