function [X, inputIsGpuArray] = moveInferenceInputToEnvironment(X, gpuShouldBeUsed, mustGatherInputsForMex)
% This function is undocumented and will change in a future release

% Ensure the input is on the selected Environment

%   Copyright 2020 The MathWorks, Inc.

arguments
    X
    gpuShouldBeUsed logical
    mustGatherInputsForMex logical = false
end

inputIsGpuArray = isgpuarray(X);

% Ensure the input is on the selected environment
if gpuShouldBeUsed
    if ~inputIsGpuArray && ~mustGatherInputsForMex
        X = gpuArray(X);
        return
    elseif inputIsGpuArray && mustGatherInputsForMex
        % MEX Acceleration may require the inputs to be host arrays
        X = gather(X);
        return
    end
end

% Error if a gpuArray is provided with CPU ExecutionEnvironment
if inputIsGpuArray && ~gpuShouldBeUsed
    exception = MException(message('nnet_cnn:DAGNetwork:InvalidDataForHostEnvironment'));
    throwAsCaller(exception);
end
end