MCPcopy Create free account
hub / github.com/deepspeedai/DeepSpeed / sequence_dataloader

Function sequence_dataloader

tests/unit/simple_model.py:282–288  ·  view source on GitHub ↗
(model, total_samples, hidden_dim, device, seq_len: int = 32, dtype=preferred_dtype())

Source from the content-addressed store, hash-verified

280
281
282def sequence_dataloader(model, total_samples, hidden_dim, device, seq_len: int = 32, dtype=preferred_dtype()):
283 batch_size = model.train_micro_batch_size_per_gpu()
284 train_data = torch.randn(total_samples, seq_len, hidden_dim, device=device, dtype=dtype)
285 train_label = torch.empty(total_samples, dtype=torch.long, device=device).random_(hidden_dim)
286 train_dataset = torch.utils.data.TensorDataset(train_data, train_label)
287 train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size)
288 return train_loader
289
290
291def create_config_from_dict(tmpdir, config_dict):

Callers 9

_train_moeFunction · 0.90
testMethod · 0.90
testMethod · 0.90
testMethod · 0.90
testMethod · 0.90
test_unfused_gradnormMethod · 0.90
test_fused_gradnormMethod · 0.90
test_lamb_gradnormMethod · 0.90

Calls 2

preferred_dtypeFunction · 0.85

Tested by 8

_train_moeFunction · 0.72
testMethod · 0.72
testMethod · 0.72
testMethod · 0.72
testMethod · 0.72
test_unfused_gradnormMethod · 0.72
test_fused_gradnormMethod · 0.72
test_lamb_gradnormMethod · 0.72