Update modeling_cxrbert.py
Browse files- modeling_cxrbert.py +3 -1
modeling_cxrbert.py
CHANGED
@@ -2,7 +2,7 @@
|
|
2 |
# Copyright (c) Microsoft Corporation. All rights reserved.
|
3 |
# Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
|
4 |
# ------------------------------------------------------------------------------------------
|
5 |
-
|
6 |
from typing import Any, Optional, Tuple, Union
|
7 |
|
8 |
import torch
|
@@ -16,6 +16,8 @@ from .configuration_cxrbert import CXRBertConfig
|
|
16 |
|
17 |
BERTTupleOutput = Tuple[T, T, T, T, T]
|
18 |
|
|
|
|
|
19 |
class CXRBertOutput(ModelOutput):
|
20 |
last_hidden_state: torch.FloatTensor
|
21 |
logits: torch.FloatTensor
|
|
|
2 |
# Copyright (c) Microsoft Corporation. All rights reserved.
|
3 |
# Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
|
4 |
# ------------------------------------------------------------------------------------------
|
5 |
+
from dataclasses import dataclass
|
6 |
from typing import Any, Optional, Tuple, Union
|
7 |
|
8 |
import torch
|
|
|
16 |
|
17 |
BERTTupleOutput = Tuple[T, T, T, T, T]
|
18 |
|
19 |
+
|
20 |
+
@dataclass
|
21 |
class CXRBertOutput(ModelOutput):
|
22 |
last_hidden_state: torch.FloatTensor
|
23 |
logits: torch.FloatTensor
|