Update README.md

README.md

@@ -101,35 +101,26 @@ for n, m in block.named_modules():
   if isinstance(m, (torch.nn.Linear, transformers.modeling_utils.Conv1D)):
     if "experts" in n and ("shared_experts" not in n) and int(n.split('.')[-2])<63  and "down_proj" not in n :
       device ="cuda:1"
-      output_device = "cuda:1"
     elif "experts" in n and ("shared_experts" not in n) and "down_proj" in n and int(n.split('.')[-2])<63:
       device = "cuda:1"
-      output_device = "cuda:0"
     elif "experts" in n and ("shared_experts" not in n) and int(n.split('.')[-2]) >= 63 and  int(n.split('.')[-2]) < 128 and "down_proj" not in n:
       device = "cuda:2"
-      output_device = "cuda:2"
     elif "experts" in n and ("shared_experts" not in n) and "down_proj" in n and int(n.split('.')[-2]) >= 63 and  int(n.split('.')[-2]) < 128:
       device = "cuda:2"
-      output_device = "cuda:0"
     elif "experts" in n and ("shared_experts" not in n) and int(n.split('.')[-2]) >= 128 and int(
           n.split('.')[-2]) < 192 and "down_proj" not in n:
       device = "cuda:3"
-      output_device = "cuda:3"
     elif "experts" in n and ("shared_experts" not in n) and "down_proj"  in n and int(
           n.split('.')[-2]) >= 128 and int(n.split('.')[-2]) < 192:
       device = "cuda:3"
-      output_device = "cuda:0"
     elif "experts" in n and ("shared_experts" not in n) and "down_proj" not in n and int(
           n.split('.')[-2]) >= 192:
       device = "cuda:4"
-      output_device = "cuda:4"
     elif "experts" in n and ("shared_experts" not in n) and "down_proj" in n and int(
           n.split('.')[-2]) >= 192:
       device = "cuda:4"
-      output_device = "cuda:0"
     else:
       device = "cuda:0"
-      output_device = "cuda:0"
     n = n[2:]
     device_map.update({n: device})