Files
PINA/pina/model/multi_feed_forward.py
2023-11-17 09:51:29 +01:00

25 lines
728 B
Python

"""Module for Multi FeedForward model"""
import torch
from .feed_forward import FeedForward
class MultiFeedForward(torch.nn.Module):
"""
The PINA implementation of MultiFeedForward network.
This model allows to create a network with multiple FeedForward combined
together. The user has to define the `forward` method choosing how to
combine the different FeedForward networks.
:param dict dff_dict: dictionary of FeedForward networks.
"""
def __init__(self, ffn_dict):
super().__init__()
if not isinstance(ffn_dict, dict):
raise TypeError
for name, constructor_args in ffn_dict.items():
setattr(self, name, FeedForward(**constructor_args))