X-Git-Url: https://www.fleuret.org/cgi-bin/gitweb/gitweb.cgi?p=pytorch.git;a=blobdiff_plain;f=conv_chain.py;h=a1d9af0fb4da60109bb0419e0566a85234169861;hp=3077874c1ea50310452bf15dd6b2f496a705eef6;hb=05b9b133a45ac9bd5abe6f8b6d29095f9c82797a;hpb=ca897077ed89fbc3c7e8d812ad262146a0c72b71

diff --git a/conv_chain.py b/conv_chain.py
index 3077874..a1d9af0 100755
--- a/conv_chain.py
+++ b/conv_chain.py
@@ -7,44 +7,49 @@
 
 ######################################################################
 
+
 def conv_chain(input_size, output_size, remain_depth, cond):
     if remain_depth == 0:
         if input_size == output_size:
-            return [ [ ] ]
+            return [[]]
         else:
-            return [ ]
+            return []
     else:
-        r = [ ]
+        r = []
         for kernel_size in range(1, input_size + 1):
             for stride in range(1, input_size):
                 if cond(remain_depth, kernel_size, stride):
                     n = (input_size - kernel_size) // stride + 1
-                    if n >= output_size and (n - 1) * stride + kernel_size == input_size:
+                    if (
+                        n >= output_size
+                        and (n - 1) * stride + kernel_size == input_size
+                    ):
                         q = conv_chain(n, output_size, remain_depth - 1, cond)
-                        r += [ [ (kernel_size, stride) ] + u for u in q ]
+                        r += [[(kernel_size, stride)] + u for u in q]
         return r
 
+
 ######################################################################
 
 if __name__ == "__main__":
-
     import torch
     from torch import nn
 
     # Example
 
     c = conv_chain(
-        input_size = 64, output_size = 8,
-        remain_depth = 5,
+        input_size=64,
+        output_size=8,
+        remain_depth=5,
         # We want kernels smaller than 4, strides smaller than the
         # kernels, and strides of 1 except in the two last layers
-        cond = lambda d, k, s: k <= 4 and s <= k and (s == 1 or d <= 2)
+        cond=lambda d, k, s: k <= 4 and s <= k and (s == 1 or d <= 2),
     )
 
     x = torch.rand(1, 1, 64)
 
     for m in c:
-        model = nn.Sequential(*[ nn.Conv1d(1, 1, l[0], l[1]) for l in m ])
+        model = nn.Sequential(*[nn.Conv1d(1, 1, l[0], l[1]) for l in m])
         print(model)
         print(x.size(), model(x).size())