code: 
  
  
  import os
  import unittest
  
  from transformers import is_torch_available
  from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device
  
  
  if is_torch_available():
  		import torch
  
  		from transformers import GPTaLMHeadModel
  
  
  
  
  @require_torch
  @require_sentencepiece
  @require_tokenizers
  class 					_UpperCamelCase (      unittest.TestCase	):
  
  
  
  
  
  
  
  							'''simple docstring'''
  
  
  
  
  
  
  
  							@slow
  							@unittest.skip('Model is not available.'		)
  							def 		_UpperCAmelCase	(		self  :				int		):
  													_a     =			'nvidia/megatron-gpt2-345m/'
  													if "MYDIR" in os.environ:
  																			_a     =			os.path.join(os.environ['MYDIR']		,				SCREAMING_SNAKE_CASE_		)
  													_a     =			GPTaLMHeadModel.from_pretrained(SCREAMING_SNAKE_CASE_		)
  													model.to(SCREAMING_SNAKE_CASE_		)
  													model.half()
  
  													_a     =			torch.tensor(
  													    [[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]]		,				device=SCREAMING_SNAKE_CASE_		,				dtype=torch.long		,				)
  
  													with torch.no_grad():
  																			_a     =			model(SCREAMING_SNAKE_CASE_		).logits
  
  													_a     =			torch.Size((1, 9, 5_0_2_5_7)		)
  													self.assertEqual(output.shape		,				SCREAMING_SNAKE_CASE_		)
  
  													_a     =			torch.tensor(
  													    [
  													        4.9414,
  													        -0.2920,
  													        -1.2148,
  													        -4.0273,
  													        -0.5161,
  													        -5.2109,
  													        -1.2412,
  													        -1.8301,
  													        -1.7734,
  													        -4.7148,
  													        -0.2317,
  													        -1.0811,
  													        -2.1777,
  													        0.4141,
  													        -3.7969,
  													        -4.0586,
  													        -2.5332,
  													        -3.3809,
  													        4.3867,
  													    ]		,				device=SCREAMING_SNAKE_CASE_		,				dtype=torch.half		,				)
  
  													for i in range(1_9		):
  																			_a ,	_a     =			8 * i // 1_7, 2_7_9_2 * i  # along the diagonal
  																			_a ,	_a     =			output[0, r, c], expected_diag[i]
  																			_a     =			f"""row={r} col={c} computed={computed} expected={expected}"""
  																			self.assertAlmostEqual(SCREAMING_SNAKE_CASE_		,				SCREAMING_SNAKE_CASE_		,				delta=1e-4		,				msg=SCREAMING_SNAKE_CASE_		)
  
  
  
style_context: 
  
  
  from .chunk_utils import chunk_layer
  from .data_transforms import make_atomaa_masks
  from .feats import atomaa_to_atomaa, frames_and_literature_positions_to_atomaa_pos, torsion_angles_to_frames
  from .loss import compute_predicted_aligned_error, compute_tm
  from .protein import Protein as OFProtein
  from .protein import to_pdb
  from .rigid_utils import Rigid, Rotation
  from .tensor_utils import dict_multimap, flatten_final_dims, permute_final_dims
  
  
  