Additional notebook tests

elk/rwkv_lm/test_rwkv.ipynb

@@ -518,7 +518,7 @@
  },
  {
  "cell_type": "code",
- "execution_count": 57,
+ "execution_count": 58,
  "metadata": {},
  "outputs": [
  {
@@ -531,42 +531,53 @@
  {
  "data": {
  "text/plain": [
- "2048"
+ "torch.Size([243, 2048])"
  ]
  },
- "execution_count": 57,
+ "execution_count": 58,
  "metadata": {},
  "output_type": "execute_result"
  }
  ],
  "source": [
  "tokens = tokenizer.encode(bad_review_text, return_tensors=\"pt\")\n",
  "print(f\"Input IDs shape = {tokens.shape}\")\n",
+ "\n",
  "states = model(tokens)\n",
  "tensor_final_state = states[-1]\n",
- "tensor_final_state.shape\n",
- "len(states[0][0])"
+ "tensor_final_state.shape"
  ]
  },
  {
  "cell_type": "code",
- "execution_count": 53,
+ "execution_count": 59,
  "metadata": {},
  "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "Input IDs shape = torch.Size([1, 243])\n"
+ ]
+ },
  {
  "data": {
  "text/plain": [
- "tensor([ 0.3164, -0.3691, 0.0165, ..., -0.0081, 0.0175, -0.0771],\n",
- " device='cuda:0', dtype=torch.bfloat16)"
+ "torch.Size([2048])"
  ]
  },
- "execution_count": 53,
+ "execution_count": 59,
  "metadata": {},
  "output_type": "execute_result"
  }
  ],
  "source": [
- "tensor_final_state[-1]"
+ "tokens = tokenizer.encode(bad_review_text, return_tensors=\"pt\")\n",
+ "print(f\"Input IDs shape = {tokens.shape}\")\n",
+ "\n",
+ "states = model(tokens[0])\n",
+ "tensor_final_state = states[-1]\n",
+ "tensor_final_state.shape"
  ]
  }
  ],