fixed for dense input with dim=1

vangjiwei · Dec 8, 2020 · 6fe5478 · 6fe5478
1 parent 65d32f9
commit 6fe5478
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/layers/nnom_input.c b/src/layers/nnom_input.c
@@ -53,7 +53,9 @@ nnom_layer_t *input_s(const nnom_io_config_t* config)
 	*/
 
 	// set parameters
-	if(config->tensor->num_dim == 2) // test for 1d input, expend h = 1
+    if(config->tensor->num_dim == 1) // test for 1d input, expend h = 1
+        layer->shape = shape(1, 1, config->tensor->dim[0]);
+    else if (config->tensor->num_dim == 2) // test for 1d input, expend h = 1
 		layer->shape = shape(1, config->tensor->dim[0], config->tensor->dim[1]);
 	else
 		layer->shape = shape(config->tensor->dim[0], config->tensor->dim[1], config->tensor->dim[2]);