Unverified Commit e615cfc3 authored by アマデウス's avatar アマデウス Committed by GitHub
Browse files

[NFC] polish test component gpt code style (#1567)

parent 6159d454
......@@ -47,8 +47,15 @@ class GPTLMModel(nn.Module):
# Only return lm_logits
return self.model(input_ids=input_ids, attention_mask=attention_mask, use_cache=not self.checkpoint)[0]
def gpt2_micro(checkpoint=True):
return GPTLMModel(checkpoint=checkpoint, hidden_size=32, num_layers=2, num_attention_heads=4, max_seq_len=64, vocab_size=128)
return GPTLMModel(checkpoint=checkpoint,
hidden_size=32,
num_layers=2,
num_attention_heads=4,
max_seq_len=64,
vocab_size=128)
def gpt2_s(checkpoint=True):
return GPTLMModel(checkpoint=checkpoint)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment