'batched' function looks buggy

https://github.com/facebookresearch/dlrm/blob/b631a99cf5ff320272d18a776cda85b5207bdf19/torchrec_dlrm/dlrm_main.py#L368

I don't think batched function is doing what we want it to be doing. What is the intended output from it?

Test case:
```
import itertools
from typing import Iterator

def batched(it: Iterator, n: int):
assert n >= 1
for x in it:
yield itertools.chain((x,), itertools.islice(it, n - 1))

lis = [i for i in range(10)]
n = 5
for b in batched(lis, n):
print(b)
for item in b:
print(item)

```

Output:

0
0
1
2
3

1
0
1
2
3

2
0
1
2
3

3
0
1
2
3

4
0
1
2
3

5
0
1
2
3

6
0
1
2
3

7
0
1
2
3

8
0
1
2
3

9
0
1
2
3

Start at torchrec_dlrm/dlrm_main.py around line 368 and run the issue's batched test case to reproduce the output. Determine the intended batching behavior with the maintainers, then add a regression test and confirm the function produces the agreed output.

Written by the indexing model from the issue text.