Update docstring for multiple output heads

numenta · Jun 8, 2021 · f25baa9 · f25baa9
1 parent cf90ff6
commit f25baa9
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/nupic/research/frameworks/dendrites/modules/dendritic_mlp.py b/nupic/research/frameworks/dendrites/modules/dendritic_mlp.py
@@ -40,7 +40,9 @@ class DendriticMLP(nn.Module):
     initializations and learning parameters
 
     :param input_size: size of the input to the network
-    :param output_size: the number of units in the output layer
+    :param output_size: the number of units in the output layer. Must be either an
+                        integer if there is a single output head, or an iterable
+                        of integers if there are multiple output heads.
     :param hidden_sizes: the number of units in each hidden layer
     :param num_segments: the number of dendritic segments that each hidden unit has
     :param dim_context: the size of the context input to the network