Fixed failing tests due to change to temporal_memory.py