Skip to content

Commit

Permalink
Exclude kpt, kpts and some attributes from to_gpu function
Browse files Browse the repository at this point in the history
  • Loading branch information
sunqm committed Nov 17, 2024
1 parent c918ee1 commit fee6c15
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion pyscf/lib/misc.py
Original file line number Diff line number Diff line change
Expand Up @@ -1533,11 +1533,14 @@ def to_gpu(method, out=None):
out_keys = set(out.__dict__).union(*cls_keys)
# Only overwrite the attributes of the same name.
keys = set(method.__dict__).intersection(out_keys)
# Keys that are not required to convert to cupy array
keep_in_nparray = {'kpt', 'kpts', 'frozen'}

for key in keys:
val = getattr(method, key)
if isinstance(val, numpy.ndarray):
val = cupy.asarray(val)
if key not in keep_in_nparray:
val = cupy.asarray(val)
elif hasattr(val, 'to_gpu'):
val = val.to_gpu()
setattr(out, key, val)
Expand Down

0 comments on commit fee6c15

Please sign in to comment.