atomai
atomai copied to clipboard
3D fcnn
Updated fcnn blocks with 3D versions. I am wondering if separate 3D versions of fcnns should be created or if existing fcnns should be updated with ndim arg in fcnn.py.
Thank you, Tommy. I think 3D convolutions can be simply used as drop-ins. Can you please write it in a more compact way, without if/else? For example,
def get_conv(dim: int) -> Type[nn.Module]:
conv_dict = {1: nn.Conv1d, 2: nn.Conv2d, 3: nn.Conv3d}
return conv_dict[dim]
PS. Not sure whether List is a correct type annotation for **kwargs.
I am not sure what output type I should put for the type annotation.
I am not sure what output type I should put for the type annotation.
nn.Module