Merge pull request #167 from devrimcavusoglu/patch-1

Update utils.py
Unbabel · Oct 11, 2023 · ae37d0a · ae37d0a
2 parents 2f1d6e0 + 2a19472
commit ae37d0a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/comet/models/utils.py b/comet/models/utils.py
@@ -12,19 +12,22 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from dataclasses import dataclass
 from typing import List
 
 import torch
 from torch.utils.data import Sampler
 from transformers.utils import ModelOutput
 
 
+@dataclass
 class Prediction(ModelOutput):
     """Renamed ModelOutput"""
 
     pass
 
 
+@dataclass
 class Target(ModelOutput):
     """Renamed ModelOutput into Targets to keep same behaviour"""