Add InputDevice specifier to random number generators shape_like input.

Signed-off-by: Michał Zientkiewicz <[email protected]>
NVIDIA · Sep 18, 2024 · 040c49d · 040c49d
1 parent e476e05
commit 040c49d
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 0 deletions.
diff --git a/dali/operators/random/beta_distribution_cpu.cc b/dali/operators/random/beta_distribution_cpu.cc
@@ -38,6 +38,7 @@ a single value per sample is generated.
     .NumInput(0, 1)
     .InputDox(0, "shape_like", "TensorList",
               "Shape of this input will be used to infer the shape of the output, if provided.")
+    .InputDevice(0, InputDevice::Metadata)
     .NumOutput(1)
     .AddOptionalArg("alpha", R"code(The alpha parameter, a positive ``float32`` scalar.)code", 1.0f,
                     true)

diff --git a/dali/operators/random/choice_cpu.cc b/dali/operators/random/choice_cpu.cc
@@ -42,6 +42,7 @@ that is: :meth:`nvidia.dali.types.DALIDataType`, :meth:`nvidia.dali.types.DALIIm
               "Otherwise ``__a`` is treated as 1D array of input samples.")
     .InputDox(1, "shape_like", "TensorList",
               "Shape of this input will be used to infer the shape of the output, if provided.")
+    .InputDevice(1, InputDevice::Metadata)
     .NumOutput(1)
     .AddOptionalArg<std::vector<float>>("p",
                                         "Distribution of the probabilities. "

diff --git a/dali/operators/random/coin_flip_cpu.cc b/dali/operators/random/coin_flip_cpu.cc
@@ -30,6 +30,7 @@ a single value per sample is generated.
     .NumInput(0, 1)
     .InputDox(0, "shape_like", "TensorList",
               "Shape of this input will be used to infer the shape of the output, if provided.")
+    .InputDevice(0, InputDevice::Metadata)
     .NumOutput(1)
     .AddOptionalArg<float>("probability",
       R"code(Probability of value 1.)code",
@@ -51,6 +52,7 @@ sample is generated.
     .NumInput(0, 1)
     .InputDox(0, "shape_like", "TensorList",
               "Shape of this input will be used to infer the shape of the output, if provided.")
+    .InputDevice(0, InputDevice::Metadata)
     .NumOutput(1)
     .AddParent("random__CoinFlip")
     .Deprecate("random__CoinFlip");  // Deprecated in 0.30

diff --git a/dali/operators/random/uniform_distribution_cpu.cc b/dali/operators/random/uniform_distribution_cpu.cc
@@ -32,6 +32,7 @@ a single value per sample is generated.
     .NumInput(0, 1)
     .InputDox(0, "shape_like", "TensorList",
               "Shape of this input will be used to infer the shape of the output, if provided.")
+    .InputDevice(0, InputDevice::Metadata)
     .NumOutput(1)
     .AddOptionalArg("range",
       R"code(Range ``[min, max)`` of a continuous uniform distribution.
@@ -67,6 +68,7 @@ a single value per sample is generated.
     .NumInput(0, 1)
     .InputDox(0, "shape_like", "TensorList",
               "Shape of this input will be used to infer the shape of the output, if provided.")
+    .InputDevice(0, InputDevice::Metadata)
     .NumOutput(1)
     .AddOptionalArg("range",
       R"code(Range ``[min, max)`` of a continuous uniform distribution.