train.py (2 lines):
	- line 546: assert network_input.shape[-1] == 3072, 'TODO: support other image sizes'
	- line 725: # TODO: remove this constraint