Imported from https://github.com/DeadAt0m/ActiveSparseShifts-PyTorch