File tree Expand file tree Collapse file tree 1 file changed +6
-1
lines changed Expand file tree Collapse file tree 1 file changed +6
-1
lines changed Original file line number Diff line number Diff line change @@ -60,12 +60,17 @@ class IntermediateTensors:
6060 tensors : dict [str , torch .Tensor ]
6161 kv_connector_output : KVConnectorOutput | None
6262
63- def __init__ (self , tensors ):
63+ def __init__ (
64+ self ,
65+ tensors : dict [str , torch .Tensor ],
66+ kv_connector_output : KVConnectorOutput | None = None ,
67+ ) -> None :
6468 # manually define this function, so that
6569 # Dynamo knows `IntermediateTensors()` comes from this file.
6670 # Otherwise, dataclass will generate this function by evaluating
6771 # a string, and we will lose the information about the source file.
6872 self .tensors = tensors
73+ self .kv_connector_output = kv_connector_output
6974
7075 def __getitem__ (self , key : str | slice ):
7176 if isinstance (key , str ):
You can’t perform that action at this time.
0 commit comments