@@ -31,7 +31,7 @@ def __init__(self, args):
3131 self .context_dim = cdim = 64
3232 args .corr_levels = 4
3333 args .corr_radius = 3
34-
34+
3535 else :
3636 self .hidden_dim = hdim = 128
3737 self .context_dim = cdim = 128
@@ -46,12 +46,12 @@ def __init__(self, args):
4646
4747 # feature network, context network, and update block
4848 if args .small :
49- self .fnet = SmallEncoder (output_dim = 128 , norm_fn = 'instance' , dropout = args .dropout )
49+ self .fnet = SmallEncoder (output_dim = 128 , norm_fn = 'instance' , dropout = args .dropout )
5050 self .cnet = SmallEncoder (output_dim = hdim + cdim , norm_fn = 'none' , dropout = args .dropout )
5151 self .update_block = SmallUpdateBlock (self .args , hidden_dim = hdim )
5252
5353 else :
54- self .fnet = BasicEncoder (output_dim = 256 , norm_fn = 'instance' , dropout = args .dropout )
54+ self .fnet = BasicEncoder (output_dim = 256 , norm_fn = 'instance' , dropout = args .dropout )
5555 self .cnet = BasicEncoder (output_dim = hdim + cdim , norm_fn = 'batch' , dropout = args .dropout )
5656 self .update_block = BasicUpdateBlock (self .args , hidden_dim = hdim )
5757
@@ -97,8 +97,8 @@ def forward(self, image1, image2, iters=torch.tensor(12), flow_init=torch.tensor
9797
9898 # run the feature network
9999 with autocast (enabled = self .args .mixed_precision ):
100- fmap1 , fmap2 = self .fnet ([image1 , image2 ])
101-
100+ fmap1 , fmap2 = self .fnet ([image1 , image2 ])
101+
102102 fmap1 = fmap1 .float ()
103103 fmap2 = fmap2 .float ()
104104 if self .args .alternate_corr :
@@ -135,10 +135,10 @@ def forward(self, image1, image2, iters=torch.tensor(12), flow_init=torch.tensor
135135 flow_up = upflow8 (coords1 - coords0 )
136136 else :
137137 flow_up = self .upsample_flow (coords1 - coords0 , up_mask )
138-
138+
139139 flow_predictions .append (flow_up )
140140
141141 if test_mode :
142142 return coords1 - coords0 , flow_up
143-
143+
144144 return flow_predictions
0 commit comments