Skip to content

Commit 1a39202

Browse files
author
Joseph Watson
committed
Set a fixed cache dir (./schedules/) for diffusion schedule
1 parent 1b476ad commit 1a39202

File tree

1 file changed

+6
-2
lines changed

1 file changed

+6
-2
lines changed

inference/model_runners.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,10 @@
2727
TOR_CAN_FLIP = util.torsion_can_flip
2828
REF_ANGLES = util.reference_angles
2929

30+
# Check for cache schedule
31+
if not os.path.exists(f'{SCRIPT_DIR}/../schedules'):
32+
os.mkdir(f'{SCRIPT_DIR}/../schedules')
33+
3034
class Sampler:
3135

3236
def __init__(self, conf: DictConfig):
@@ -37,7 +41,7 @@ def __init__(self, conf: DictConfig):
3741
"""
3842
self.initialized = False
3943
self.initialize(conf)
40-
44+
4145
def initialize(self, conf: DictConfig) -> None:
4246
"""
4347
Initialize sampler.
@@ -113,7 +117,7 @@ def initialize(self, conf: DictConfig) -> None:
113117
self.potential_conf = self._conf.potentials
114118
self.diffuser_conf = self._conf.diffuser
115119
self.preprocess_conf = self._conf.preprocess
116-
self.diffuser = Diffuser(**self._conf.diffuser)
120+
self.diffuser = Diffuser(**self._conf.diffuser, cache_dir=f'{SCRIPT_DIR}/../schedules')
117121

118122
###########################
119123
### Initialise Symmetry ###

0 commit comments

Comments
 (0)