Add is_training to MLP and U-Net#11

copybara-service[bot] wants to merge 1 commit intomaingoogle/hackable_diffusion:mainfrom

test_872361800google/hackable_diffusion:test_872361800Copy head branch name to clipboard