diff --git a/R/PipeOpTaskPreprocTorch.R b/R/PipeOpTaskPreprocTorch.R
index 2d39f25f..a3f52cda 100644
--- a/R/PipeOpTaskPreprocTorch.R
+++ b/R/PipeOpTaskPreprocTorch.R
@@ -421,6 +421,7 @@ create_ps = function(fn) {
 #'   If "infer", the output shape function is inferred and calculates the output shapes as follows:
 #'   For an input shape of (NA, ...) a meta-tensor of shape (1, ...) is created and the preprocessing function is
 #'   applied. Afterwards the batch dimension (1) is replaced with NA and the shape is returned.
+#'  If the first dimension is not `NA`, the output shape of applying the preprocessing function is returned.
 #'
 #'   This should be correct in most cases, but might fail in some edge cases.
 #' @param param_set ([`ParamSet`] or `NULL`)\cr