`torch.rand` is treated as a constant by ONNX

rand.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Install required dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "m-28wslFI2wX",
        "outputId": "42aedcf7-3ae6-4df0-b891-c3c18895c883"
      },
      "outputs": [],
      "source": [
        "%pip install torch\n",
        "%pip install onnxruntime\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Model Definition"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-61wEcCP1OZ5"
      },
      "outputs": [],
      "source": [
        "import torch.nn as nn\n",
        "import torch\n",
        "class NN(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(NN, self).__init__()\n",
        "\n",
        "    def forward(self):\n",
        "        return torch.rand(torch.tensor(1), 1)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Export the torch model to ONNX"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "\n",
        "import torch\n",
        "import os\n",
        "\n",
        "\n",
        "out_dir = 'model'\n",
        "os.makedirs(out_dir, exist_ok=True)\n",
        "model = NN()\n",
        "onnx_path = os.path.join(out_dir, \"network.onnx\")\n",
        "\n",
        "# Neither eval nor compile effect the issue.\n",
        "model.eval()\n",
        "#torch.compile(model)\n",
        "\n",
        "torch.onnx.export(\n",
        "                model,\n",
        "                (),\n",
        "                onnx_path,\n",
        "                export_params=True,\n",
        "                opset_version=14, # Opset version doesn't seem to have an effect on the issue.\n",
        "                do_constant_folding=False, # Constant folding has no effect on the issue.\n",
        "                input_names=[],\n",
        "                output_names=['output'],\n",
        "            )\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Reproduction of the issue\n",
        "\n",
        "Despite the forward method containing `torch.rand()`, the ONNX treats the random value as a constant."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "import onnxruntime as ort\n",
        "\n",
        "session = ort.InferenceSession(onnx_path)\n",
        "results = session.run([], {})\n",
        "print(\"ONNX Raw: \", results)\n",
        "torch_out = model()\n",
        "print(\"Torch Raw: \", torch_out)\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}