File size: 433 Bytes
78a947a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
"""
Contains an implementation of a Module that sequentially applies multiple models.
Author: Ole-Christian Galbo Engstrøm
E-mail: ocge@foss.dk
"""
import torch.nn as nn
class CompoundModel(nn.Module):
def __init__(self, *models):
super(CompoundModel, self).__init__()
self.models = nn.ModuleList(models)
def forward(self, x):
for model in self.models:
x = model(x)
return x
|