Unverified Commit a0531c8a authored by Michael Benayoun's avatar Michael Benayoun Committed by GitHub
Browse files

fixed shape issue for T5 tracing (#11742)


Co-authored-by: default avatarMichael Benayoun <michael@huggingface.co>
parent 0fc56df5
...@@ -68,6 +68,8 @@ class HFProxy(Proxy): ...@@ -68,6 +68,8 @@ class HFProxy(Proxy):
if self.tracer.num_choices <= 0: if self.tracer.num_choices <= 0:
raise ValueError("num_choices must be given to the CustomTracer for MultipleChoice tasks.") raise ValueError("num_choices must be given to the CustomTracer for MultipleChoice tasks.")
shape = shape[:1] + [self.tracer.num_choices] + shape[1:] shape = shape[:1] + [self.tracer.num_choices] + shape[1:]
elif "hidden_states.s" in code_context:
shape = shape + [self.tracer.root.config.hidden_size]
else: else:
# Default case: # Default case:
# - If self.size is called for an unpacking, retrieves the corresponding unpacking # - If self.size is called for an unpacking, retrieves the corresponding unpacking
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment