From 084fc89d634350123337b71d62690a2f69093b90 Mon Sep 17 00:00:00 2001
From: Shiliang Chen <36809537+csl122@users.noreply.github.com>
Date: Mon, 8 May 2023 16:50:28 +0100
Subject: [PATCH] add two missing param comments for UNetModel class

comments for tf_layers and d_cond were missing, added.
---
 labml_nn/diffusion/stable_diffusion/model/unet.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/labml_nn/diffusion/stable_diffusion/model/unet.py b/labml_nn/diffusion/stable_diffusion/model/unet.py
index 311482dd..761bded9 100644
--- a/labml_nn/diffusion/stable_diffusion/model/unet.py
+++ b/labml_nn/diffusion/stable_diffusion/model/unet.py
@@ -49,7 +49,9 @@ def __init__(
         :param n_res_blocks: number of residual blocks at each level
         :param attention_levels: are the levels at which attention should be performed
         :param channel_multipliers: are the multiplicative factors for number of channels for each level
-        :param n_heads: the number of attention heads in the transformers
+        :param n_heads: is the number of attention heads in the transformers
+        :param tf_layers: is the number of transformer layers in the transformers
+        :param d_cond: is the size of the conditional embedding in the transformers
         """
         super().__init__()
         self.channels = channels