nvidia/C-RADIO

1# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.2#3# NVIDIA CORPORATION and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto.  Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION is strictly prohibited.8from argparse import Namespace9from typing import NamedTuple10 11import torch12from torch import nn13import torch.nn.functional as F14 15 16class AdaptorInput(NamedTuple):17    images: torch.Tensor18    summary: torch.Tensor19    features: torch.Tensor20 21 22class RadioOutput(NamedTuple):23    summary: torch.Tensor24    features: torch.Tensor25 26    def to(self, *args, **kwargs):27        return RadioOutput(28            self.summary.to(*args, **kwargs) if self.summary is not None else None,29            self.features.to(*args, **kwargs) if self.features is not None else None,30        )31 32 33class AdaptorBase(nn.Module):34    def forward(self, input: AdaptorInput) -> RadioOutput:35        raise NotImplementedError("Subclasses must implement this!")36