Modify get_flops.py for editing tasks
@@ -26,6 +26,8 @@ def main():
input_shape = (3, args.shape[0], args.shape[0])
elif len(args.shape) == 2:
input_shape = (3, ) + tuple(args.shape)
+ elif len(args.shape) == 3:
+ input_shape = tuple(args.shape)
else:
raise ValueError('invalid input shape')
