From ae666dc520b3117da55d643700cde3a988d267b4 Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Sat, 19 Dec 2020 10:28:23 -0700
Subject: [PATCH] Fix bugs with srflow after refactor

---
 codes/models/srflow/FlowStep.py       | 2 ++
 recipes/srflow/train_div2k_srflow.yml | 4 ++--
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/codes/models/srflow/FlowStep.py b/codes/models/srflow/FlowStep.py
index 9f6e4f8d..b62d4b55 100644
--- a/codes/models/srflow/FlowStep.py
+++ b/codes/models/srflow/FlowStep.py
@@ -2,6 +2,8 @@ import torch
 from torch import nn as nn
 
 import models.srflow.Permutations
+import models.srflow.FlowAffineCouplingsAblation
+import models.srflow.FlowActNorms
 
 
 def getConditional(rrdbResults, position):
diff --git a/recipes/srflow/train_div2k_srflow.yml b/recipes/srflow/train_div2k_srflow.yml
index 3da10899..122576a6 100644
--- a/recipes/srflow/train_div2k_srflow.yml
+++ b/recipes/srflow/train_div2k_srflow.yml
@@ -24,12 +24,12 @@ datasets:
 networks:
   generator:
     type: generator
-    which_model_G: srflow_orig
+    which_model_G: srflow
     nf: 64
     nb: 23
     K: 16
     scale: 4
-    initial_stride: 2
+    initial_stride: 1
     flow_scale: 4
     train_RRDB: false    # <-- Start false. After some time, ~20k-50k steps, set to true. TODO: automate this.
     pretrain_rrdb: ../experiments/pretrained_rrdb.pth  # <-- Insert path to your pretrained RRDB here.