Specifically address https://github.com/AllenCellModeling/pytorch_fnet/pull/109#discussion_r275568478