wip

Author: tastelikefeet

src/twinkle/data_format/input_feature.py

@@ -13,7 +13,7 @@
 if TYPE_CHECKING:
     import torch
 
-InputType = Union[List[List[int]], List[int], np.ndarray, torch.Tensor]
+InputType = Union[List[List[int]], List[int], np.ndarray, 'torch.Tensor']
 
 
 class InputFeature(TypedDict, total=False):

src/twinkle/data_format/output.py

@@ -12,7 +12,7 @@
 if TYPE_CHECKING:
     import torch
 
-OutputType = Union[np.ndarray, torch.Tensor, List[Any]]
+OutputType = Union[np.ndarray, 'torch.Tensor', List[Any]]
 
 
 class ModelOutput(TypedDict, total=False):

src/twinkle/model/base.py

@@ -60,15 +60,15 @@ def clip_grad_and_step(self, max_grad_norm: float=1.0, norm_type=2, **kwargs):
         ...
 
     @abstractmethod
-    def set_loss(self, loss_cls: Union[Loss, Type[Loss], str, Callable[[InputFeature, ModelOutput, ...], torch.Tensor]], **kwargs):
+    def set_loss(self, loss_cls: Union[Loss, Type[Loss], str, Callable[[InputFeature, ModelOutput, ...], 'torch.Tensor']], **kwargs):
         ...
 
     @abstractmethod
-    def set_optimizer(self, optimizer_cls: Union[Optimizer, Type[Optimizer], str], **kwargs):
+    def set_optimizer(self, optimizer_cls: Union['Optimizer', Type['Optimizer'], str], **kwargs):
         ...
 
     @abstractmethod
-    def set_lr_scheduler(self, scheduler_cls: Union[LRScheduler, Type[LRScheduler], str], **kwargs):
+    def set_lr_scheduler(self, scheduler_cls: Union['LRScheduler', Type['LRScheduler'], str], **kwargs):
         ...
 
     @abstractmethod

src/twinkle/model/transformers/transformers.py

@@ -384,6 +384,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
             inputs['labels'] = labels
         optimizer_config.inputs = inputs
         optimizer_config.outputs = outputs
+        optimizer_config.loss_value = outputs.get('aux_loss', 0)
         return outputs
 
     @remote_function(collect='mean')