From 04bdd7bbf237f76282ca0851c36461bb87148438 Mon Sep 17 00:00:00 2001 From: "Dr. Kashif Rasul" Date: Fri, 17 Jan 2020 11:42:46 +0100 Subject: [PATCH] added MAF option to model --- pts/model/tempflow/tempflow_network.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pts/model/tempflow/tempflow_network.py b/pts/model/tempflow/tempflow_network.py index bbb9c93..92bc2af 100644 --- a/pts/model/tempflow/tempflow_network.py +++ b/pts/model/tempflow/tempflow_network.py @@ -5,7 +5,7 @@ import torch.nn as nn import numpy as np -from pts.modules import RealNVP, FlowOutput, MeanScaler, NOPScaler +from pts.modules import RealNVP, MAF, FlowOutput, MeanScaler, NOPScaler from pts.model import weighted_average @@ -54,7 +54,7 @@ class TempFlowTrainingNetwork(nn.Module): batch_first=True, ) - flow_cls = {"RealNVP": RealNVP}[flow_type] + flow_cls = {"RealNVP": RealNVP, "MAF": MAF,}[flow_type] self.flow = flow_cls( input_size=target_dim, n_blocks=n_blocks,