Skip to content

Commit

Permalink
adds HalfUnet to illustrate the doc
Browse files Browse the repository at this point in the history
  • Loading branch information
Frank Guibert committed Jul 4, 2024
1 parent 89a6cd2 commit 4ee1fec
Showing 1 changed file with 10 additions and 0 deletions.
10 changes: 10 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,16 @@ Each model we provide is a subclass of [torch.nn.Module](https://pytorch.org/doc
- **onnx_supported**: a boolean that indicates if the model can be exported to onnx. Our CI validates that the model can be exported to onnx and reloaded for inference.

```python
@dataclass_json
@dataclass(slots=True)
class HalfUNetSettings:
num_filters: int = 64
dilation: int = 1
bias: bool = False
use_ghost: bool = False
last_activation: str = "Identity"
absolute_pos_embed: bool = False

class HalfUNet(nn.Module):
settings_kls = HalfUNetSettings
onnx_supported = True
Expand Down

0 comments on commit 4ee1fec

Please sign in to comment.