YaohuiW commited on
Commit
3647938
·
verified ·
1 Parent(s): ac6b29b

Update networks/generator.py

Browse files
Files changed (1) hide show
  1. networks/generator.py +1 -0
networks/generator.py CHANGED
@@ -4,6 +4,7 @@ from networks.encoder import Encoder
4
  from networks.decoder import Decoder
5
  import numpy as np
6
  from tqdm import tqdm
 
7
 
8
 
9
  class Generator(nn.Module):
 
4
  from networks.decoder import Decoder
5
  import numpy as np
6
  from tqdm import tqdm
7
+ from einops import rearrange
8
 
9
 
10
  class Generator(nn.Module):