Fixing build and doc errors after PR merge (#154)

* update docs references and remove implicit optional in unet.py * trying PEP604 syntax
pinellolab · Jun 23, 2023 · a55002d · a55002d
1 parent d93906b
commit a55002d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 19 deletions.
diff --git a/docs/reference/dnadiffusion.md b/docs/reference/dnadiffusion.md
@@ -2,17 +2,7 @@
 
 ## dnadiffusion
 
-::: dnadiffusion.trainer
-    options:
-      show_root_heading: true
-      show_source: false
-
-::: dnadiffusion.configs
-    options:
-      show_root_heading: true
-      show_source: false
-
-::: dnadiffusion.sample
+::: dnadiffusion.utils.train_util
     options:
       show_root_heading: true
       show_source: false
@@ -26,17 +16,12 @@
 
 ## models
 
-::: dnadiffusion.models.modules
+::: dnadiffusion.models.layers
     options:
       show_root_heading: true
       show_source: false
 
-::: dnadiffusion.models.training_modules
-    options:
-      show_root_heading: true
-      show_source: false
-
-::: dnadiffusion.models.networks
+::: dnadiffusion.models.diffusion
     options:
       show_root_heading: true
       show_source: false

diff --git a/src/dnadiffusion/models/unet.py b/src/dnadiffusion/models/unet.py
@@ -1,4 +1,5 @@
 from functools import partial
+from typing import Optional
 
 from memory_efficient_attention_pytorch import Attention as EfficientAttention
 
@@ -9,7 +10,7 @@ class UNet(nn.Module):
     def __init__(
         self,
         dim: int,
-        init_dim: Optional[int] = None,
+        init_dim: int | None = None,
         dim_mults: tuple = (1, 2, 4),
         channels: int = 1,
         resnet_block_groups: int = 8,