padding not needed, as not doing ring passing

assert_zig_zag.py

@@ -92,23 +92,19 @@ def start(
 
     # zig zag
 
-    padded_zig_zag_input, remove_pad = zig_zag_pad_seq(zig_zag_input)
+    inp, gather_seq = zig_zag_shard(zig_zag_input, all_gather_batch = True)
 
-    padded_zig_zag_input, gather_seq = zig_zag_shard(padded_zig_zag_input, all_gather_batch = True)
-
-    qkv = attention.to_qkv(padded_zig_zag_input)
+    qkv = attention.to_qkv(inp)
 
     q, k, v = rearrange(qkv, 'b n (h d) -> b h n d', d = dim_head).split(attention.qkv_head_breakdown, dim = -3)
 
     o = zig_zag_attn(q, k, v)
 
     o = rearrange(o, 'b h n d -> b n (h d)')
 
-    padded_zig_zag_out = attention.to_out(o)
-
-    padded_zig_zag_out = gather_seq(padded_zig_zag_out)
+    zig_zag_out = attention.to_out(o)
 
-    zig_zag_out = remove_pad(padded_zig_zag_out)
+    zig_zag_out = gather_seq(zig_zag_out)
 
     zig_zag_out.mean().backward()