Update README.md

README.md

@@ -45,7 +45,7 @@ Mimic human preference given a query and 2 different answers.
 import torch
 from torch import nn
 from torch.nn import CrossEntropyLoss, MSELoss, BCEWithLogitsLoss
-from transformers import Gemma2PreTrainedModel,Gemma2Model
+from transformers import Gemma2PreTrainedModel,Gemma2Model, Cache, AutoTokenizer
 from transformers.modeling_outputs import SequenceClassifierOutputWithPast
 from typing import Optional, List, Union, Tuple
 from dataclasses import dataclass