Improve weight tying handling

Author: rasbt

ch05/07_gpt_to_llama/converting-gpt-to-llama2.ipynb

@@ -1400,14 +1400,17 @@
    },
    "outputs": [],
    "source": [
-    "def assign(left, right):\n",
+    "def assign(left, right, tensor_name=\"unknown\"):\n",
     "    if left.shape != right.shape:\n",
-    "        raise ValueError(f\"Shape mismatch. Left: {left.shape}, Right: {right.shape}\")\n",
+    "        raise ValueError(f\"Shape mismatch in tensor '{tensor_name}'. Left: {left.shape}, Right: {right.shape}\")\n",
+    "    \n",
+    "    with torch.no_grad():\n",
+    "        if isinstance(right, torch.Tensor):\n",
+    "            left.copy_(right)\n",
+    "        else:\n",
+    "            left.copy_(torch.as_tensor(right, dtype=left.dtype, device=left.device))\n",
     "\n",
-    "    if isinstance(right, torch.Tensor):\n",
-    "        return torch.nn.Parameter(right.clone().detach())\n",
-    "    else:\n",
-    "        return torch.nn.Parameter(torch.tensor(right))\n",
+    "    return left \n",
     "\n",
     "\n",
     "def permute(w: torch.Tensor, n_heads, out_dim, in_dim):\n",