{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "422VxNWK9WM8"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class NeuralNetwork(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(NeuralNetwork, self).__init__()\n",
        "\n",
        "        self.fc1 = nn.Linear(784, 256)\n",
        "        self.fc2 = nn.Linear(256, 128)\n",
        "        self.fc3 = nn.Linear(128, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = torch.relu(self.fc1(x))\n",
        "        x = torch.relu(self.fc2(x))\n",
        "        x = self.fc3(x)\n",
        "        return x"
      ],
      "metadata": {
        "id": "dH9PX_l49W6s"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = NeuralNetwork()\n",
        "\n",
        "for layer in model.modules():\n",
        "    if isinstance(layer, nn.Linear):\n",
        "        nn.init.kaiming_normal_(layer.weight, nonlinearity='relu')"
      ],
      "metadata": {
        "id": "N5jQ2kjU9W8_"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "criterion = nn.CrossEntropyLoss()\n",
        "optimizer = optim.Adam(model.parameters(), lr=0.001)"
      ],
      "metadata": {
        "id": "b77eYxVy9W_Q"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "X = torch.randn(64, 784)\n",
        "y = torch.randint(0, 10, (64,))"
      ],
      "metadata": {
        "id": "8iO2cWYR9XBf"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "outputs = model(X)\n",
        "loss = criterion(outputs, y)"
      ],
      "metadata": {
        "id": "C5nc7yv19XD3"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "optimizer.zero_grad()\n",
        "loss.backward()\n",
        "optimizer.step()"
      ],
      "metadata": {
        "id": "nXRAuDq-9XGX"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Loss:\", loss.item())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "iTy0IQAj9XIw",
        "outputId": "3aafa77b-0bec-400b-b8c5-fa615c0dcaf5"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loss: 3.227264642715454\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "Zc1ePd519XLI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "eSIQB6RI9XPX"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}