Skip to content

Commit 3dcb9a2

Browse files
orionarcheresoteric-ephemera
authored andcommitted
Fix BinningAutoBatcher/InFlightAutoBatcher ignoring model.memory_scales_with
When constructing autobatchers from a user-supplied dict, propagate model.memory_scales_with as the default so fixed-neighbor models like FairChem use the correct scaling metric instead of the constructor default.
1 parent 4360a7a commit 3dcb9a2

1 file changed

Lines changed: 2 additions & 0 deletions

File tree

src/atomate2/torchsim/core.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -244,6 +244,7 @@ def process_in_flight_autobatcher_dict(
244244
state, model, autobatcher=autobatcher_dict, max_iterations=max_iterations
245245
)
246246
else:
247+
autobatcher_dict.setdefault("memory_scales_with", model.memory_scales_with)
247248
autobatcher = InFlightAutoBatcher(model=model, **autobatcher_dict)
248249

249250
autobatcher_details = _get_autobatcher_details(autobatcher)
@@ -284,6 +285,7 @@ def process_binning_autobatcher_dict(
284285
# pop max_iterations if present
285286
autobatcher_dict = deepcopy(autobatcher_dict)
286287
autobatcher_dict.pop("max_iterations", None)
288+
autobatcher_dict.setdefault("memory_scales_with", model.memory_scales_with)
287289
autobatcher = BinningAutoBatcher(model=model, **autobatcher_dict)
288290

289291
autobatcher_details = _get_autobatcher_details(autobatcher)

0 commit comments

Comments
 (0)