Change all instances of .flat to .flatten(), which caused a problem in forward runs.
Change all instances of .flat to .flatten(), which caused a problem in forward runs.
Change all instances of .flat to .flatten(), which caused a problem in forward runs.