Add example

openclimatefix · Feb 9, 2022 · c7c651c · c7c651c
1 parent 899b21d
commit c7c651c
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -49,3 +49,22 @@ discriminator = Discriminator().from_pretrained("openclimagefix/dgmr-discriminat
 latent_stack = LatentConditioningStack().from_pretrained("openclimatefix/dgmr-latent-conditioning-stack")
 context_stack = ContextConditioningStack().from_pretrained("openclimatefix/dgmr-context-conditioning-stack")
 ```
+
+## Example Usage
+
+```python
+from dgmr import DGMR
+model = DGMR(
+        forecast_steps=4,
+        input_channels=1,
+        output_shape=128,
+        latent_channels=384,
+        context_channels=192,
+        num_samples=3,
+    )
+x = torch.rand((2, 4, 1, 128, 128))
+out = model(x)
+y = torch.rand((2, 4, 1, 128, 128))
+loss = F.mse_loss(y, out)
+loss.backward()
+```