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 Namespace9 10import torch11from torch import nn12import torch.nn.functional as F13 14from .adaptor_base import AdaptorBase, AdaptorInput, RadioOutput15from .adaptor_mlp import create_mlp_from_state16 17 18class GenericAdaptor(AdaptorBase):19    def __init__(self, main_config: Namespace, adaptor_config, state):20        super().__init__()21 22        self.head_mlp = create_mlp_from_state(main_config.mlp_version, state, 'summary.')23        self.feat_mlp = create_mlp_from_state(main_config.mlp_version, state, 'feature.')24 25    def forward(self, input: AdaptorInput) -> RadioOutput:26        summary = self.head_mlp(input.summary)27        feat = self.feat_mlp(input.features)28 29        return RadioOutput(summary, feat)30