Skip to content

Commit 2a6b375

Browse files
committed
Replace itertools.tee with list
1 parent 2a14d00 commit 2a6b375

File tree

1 file changed

+2
-5
lines changed

1 file changed

+2
-5
lines changed

torch/optim/sparse_adam.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import itertools
21
import math
32
import torch
43
from .optimizer import Optimizer
@@ -33,12 +32,10 @@ def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8):
3332
if not 0.0 <= betas[1] < 1.0:
3433
raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
3534

36-
# if params are in the form of a generator, the next for-loop exhausts it,
37-
# so the copy is passed to the loop
38-
params, params_copy = itertools.tee(params)
35+
params = list(params)
3936

4037
sparse_params = []
41-
for index, param in enumerate(params_copy):
38+
for index, param in enumerate(params):
4239
if isinstance(param, dict):
4340
for d_index, d_param in enumerate(param.get("params", [])):
4441
if d_param.is_sparse:

0 commit comments

Comments
 (0)