Runtime Error when saving model

Hello, I’m having the following run-time error when saving my model at the first model save point. Any ideas or help would be excellent. Thank you.

RuntimeError: The size of tensor a (25) must match the size of tensor b (50) at non-singleton dimension 0

More context:

Saving weights to weights/BigGAN_C100_seed1_Gch64_Dch64_bs50_nDs4_Glr2.0e-04_Dlr2.0e-04_Gnlrelu_Dnlrelu_GinitN02_DinitN02_ema/copy0...
Traceback (most recent call last):
  File "train.py", line 227, in <module>
    main()
  File "train.py", line 224, in main
    run(config)
  File "train.py", line 206, in run
    state_dict, config, experiment_name)
  File "/workspace/BigGAN-PyTorch/train_fns.py", line 140, in save_and_sample
    z_=z_)
  File "/workspace/BigGAN-PyTorch/utils.py", line 895, in sample_sheet
    o = nn.parallel.data_parallel(G, (z_[:classes_per_sheet], G.shared(y)))
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py", line 207, in data_parallel
    outputs = parallel_apply(replicas, inputs, module_kwargs, used_device_ids)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py", line 83, in parallel_apply
    raise output
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py", line 59, in _worker
    output = module(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 491, in __call__
    result = self.forward(*input, **kwargs)
  File "/workspace/BigGAN-PyTorch/BigGAN.py", line 248, in forward
    h = block(h, ys[index])
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 491, in __call__
    result = self.forward(*input, **kwargs)
  File "/workspace/BigGAN-PyTorch/layers.py", line 399, in forward
    h = self.activation(self.bn1(x, y))
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 491, in __call__
    result = self.forward(*input, **kwargs)
  File "/workspace/BigGAN-PyTorch/layers.py", line 325, in forward
    return out * gain + bias
r```

[ajbrock]

Actually one more mistake in my logic (I flipped the greater than the wrong way), but NOW this one should handle it.

[ajbrock]

There’s just an issue in my logic (it’s 10PM on a Friday night and I’m coding over an ssh connection on my phone <_<) this most recent commit should handle it.